Skip to content

Commit 6bd398d

Browse files
authored
Describe cross-component tied weights (#747)
## Summary - expose typed cross-component shared-weight metadata through the component manifest and `inspect_components()` - infer tied embedding/LM-head endpoints only when explicit source-path aliases identify a unique pair; for unquantized models the nested text tie flag takes precedence over the parent, while quantization metadata can explicitly restore a tie - materialize canonical packed sidecars into split component aliases before model-specific preprocessing ## Validation - targeted component-inspection, Gemma4, and weight-adapter suite (163 passing), Ruff and MyPy - companion microsoft/Olive#2684 adds a local tiny-Gemma4 Olive assembly -> Mobius ONNX embedding/decoder -> ORT CPU full-logit parity test Companion: microsoft/Olive#2684 handles component quantization and assembly. --------- Signed-off-by: Xiaoyu Zhang <xiaoyuzhang@microsoft.com>
1 parent 783e03b commit 6bd398d

7 files changed

Lines changed: 495 additions & 7 deletions

File tree

‎src/mobius/__init__.py‎

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,8 @@
2222
"CausalLMConfig",
2323
"CausalLMTask",
2424
"ComponentInfo",
25+
"SharedWeightEndpoint",
26+
"SharedWeightInfo",
2527
"ComponentExportDisposition",
2628
"ComponentExportReport",
2729
"DepthAnythingConfig",
@@ -113,7 +115,12 @@
113115
from mobius._constants import OPSET_VERSION
114116
from mobius._execution_providers import EpCapabilities, ep_registry, get_ep, register_ep
115117
from mobius._export_report import ComponentExportDisposition, ComponentExportReport
116-
from mobius._inspect import ComponentInfo, inspect_components
118+
from mobius._inspect import (
119+
ComponentInfo,
120+
SharedWeightEndpoint,
121+
SharedWeightInfo,
122+
inspect_components,
123+
)
117124
from mobius._model_package import ModelPackage
118125
from mobius._optimizations import optimize_model
119126
from mobius._registry import (

‎src/mobius/_component_manifest.py‎

Lines changed: 176 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,8 @@
88
__all__ = [
99
"ComponentDescriptor",
1010
"ComponentManifest",
11+
"SharedWeightEndpoint",
12+
"SharedWeightInfo",
1113
"get_hf_component_sources",
1214
"resolve_component_manifest",
1315
]
@@ -20,6 +22,61 @@
2022
if TYPE_CHECKING:
2123
pass
2224

25+
_INPUT_EMBEDDING_MODULE_NAMES = {
26+
"codec_embedding",
27+
"embed_tokens",
28+
"shared",
29+
"text_embedding",
30+
"tok_embeddings",
31+
}
32+
_OUTPUT_HEAD_MODULE_NAMES = {
33+
"codec_head",
34+
"lm_head",
35+
"output",
36+
"output_projection",
37+
"proj_out",
38+
}
39+
40+
41+
@dataclasses.dataclass(frozen=True)
42+
class SharedWeightEndpoint:
43+
"""One component-local consumer of a shared HuggingFace parameter."""
44+
45+
component: str
46+
parameter: str
47+
48+
def __post_init__(self) -> None:
49+
if not self.component:
50+
raise ValueError("shared-weight component must not be empty")
51+
if not self.parameter:
52+
raise ValueError("shared-weight parameter must not be empty")
53+
54+
55+
@dataclasses.dataclass(frozen=True)
56+
class SharedWeightInfo:
57+
"""A logical HuggingFace parameter consumed by multiple components."""
58+
59+
name: str
60+
canonical: SharedWeightEndpoint
61+
aliases: tuple[SharedWeightEndpoint, ...]
62+
kind: str = "parameter_alias"
63+
64+
def __post_init__(self) -> None:
65+
if not self.name:
66+
raise ValueError("shared-weight name must not be empty")
67+
if not self.kind:
68+
raise ValueError(f"shared weight {self.name!r} kind must not be empty")
69+
if not self.aliases:
70+
raise ValueError(f"shared weight {self.name!r} must declare at least one alias")
71+
endpoints = (self.canonical, *self.aliases)
72+
if len(set(endpoints)) != len(endpoints):
73+
raise ValueError(f"shared weight {self.name!r} contains duplicate endpoints")
74+
75+
@property
76+
def endpoints(self) -> tuple[SharedWeightEndpoint, ...]:
77+
"""Canonical endpoint followed by every alias."""
78+
return (self.canonical, *self.aliases)
79+
2380

2481
@dataclasses.dataclass(frozen=True)
2582
class ComponentDescriptor:
@@ -106,6 +163,7 @@ class ComponentManifest(Mapping[str, ComponentDescriptor]):
106163
"""Ordered, immutable component metadata keyed by package component name."""
107164

108165
components: tuple[ComponentDescriptor, ...]
166+
shared_weights: tuple[SharedWeightInfo, ...] = ()
109167
_by_name: Mapping[str, ComponentDescriptor] = dataclasses.field(
110168
init=False,
111169
repr=False,
@@ -120,6 +178,30 @@ def __post_init__(self) -> None:
120178
f"component manifest declares {component.name!r} more than once"
121179
)
122180
by_name[component.name] = component
181+
shared_names: set[str] = set()
182+
for shared_weight in self.shared_weights:
183+
if shared_weight.name in shared_names:
184+
raise ValueError(
185+
f"shared weight {shared_weight.name!r} is declared more than once"
186+
)
187+
shared_names.add(shared_weight.name)
188+
for endpoint in shared_weight.endpoints:
189+
if endpoint.component not in by_name:
190+
raise ValueError(
191+
f"shared weight {shared_weight.name!r} references unknown "
192+
f"component {endpoint.component!r}"
193+
)
194+
source_paths = by_name[endpoint.component].source_paths
195+
module_path = endpoint.parameter.rpartition(".")[0]
196+
if source_paths and not any(
197+
module_path == source_path or module_path.startswith(f"{source_path}.")
198+
for source_path in source_paths
199+
):
200+
raise ValueError(
201+
f"shared weight {shared_weight.name!r} parameter "
202+
f"{endpoint.parameter!r} is outside component "
203+
f"{endpoint.component!r} source paths {source_paths!r}"
204+
)
123205
object.__setattr__(self, "_by_name", MappingProxyType(by_name))
124206

125207
def __getitem__(self, name: str) -> ComponentDescriptor:
@@ -154,6 +236,93 @@ def get_hf_component_sources(
154236
return {name: tuple(paths) for name, paths in resolved.items()}
155237

156238

239+
def _config_ties_word_embeddings(hf_config: object) -> bool:
240+
"""Honor text-config precedence, except for an explicit quantized tie."""
241+
242+
def field(value: object, name: str) -> object | None:
243+
if isinstance(value, Mapping):
244+
return value.get(name)
245+
return getattr(value, name, None)
246+
247+
configs = (
248+
hf_config,
249+
field(hf_config, "text_config"),
250+
field(hf_config, "llm_config"),
251+
field(hf_config, "language_config"),
252+
)
253+
if any(
254+
bool(field(declaration, "tie_word_embeddings"))
255+
for config in configs
256+
if config is not None
257+
for declaration in (
258+
field(config, "quantization_config"),
259+
field(config, "quantization"),
260+
)
261+
if declaration is not None
262+
):
263+
return True
264+
text_config = next((config for config in configs[1:] if config is not None), None)
265+
for config in (text_config, hf_config):
266+
if config is None:
267+
continue
268+
for name in ("tie_word_embeddings", "weight_tying"):
269+
value = field(config, name)
270+
if value is not None:
271+
return bool(value)
272+
return False
273+
274+
275+
def _aliased_component_endpoint(
276+
manifest: ComponentManifest,
277+
*,
278+
role: str,
279+
local_names: set[str],
280+
) -> SharedWeightEndpoint | None:
281+
"""Resolve one explicitly aliased source module for a component role."""
282+
candidates = {
283+
SharedWeightEndpoint(
284+
component=component.name,
285+
parameter=f"{source_path}.weight",
286+
)
287+
for component in manifest.values()
288+
if component.role == role
289+
for local_path, source_path in component.source_path_aliases
290+
if local_path.rsplit(".", 1)[-1] in local_names
291+
}
292+
if len(candidates) != 1:
293+
return None
294+
return candidates.pop()
295+
296+
297+
def _infer_shared_weights(
298+
hf_config: object,
299+
manifest: ComponentManifest,
300+
) -> tuple[SharedWeightInfo, ...]:
301+
"""Infer tied embeddings only for explicit, unique component source aliases."""
302+
if not _config_ties_word_embeddings(hf_config):
303+
return ()
304+
canonical = _aliased_component_endpoint(
305+
manifest,
306+
role="embedding",
307+
local_names=_INPUT_EMBEDDING_MODULE_NAMES,
308+
)
309+
output = _aliased_component_endpoint(
310+
manifest,
311+
role="decoder",
312+
local_names=_OUTPUT_HEAD_MODULE_NAMES,
313+
)
314+
if canonical is None or output is None or canonical.component == output.component:
315+
return ()
316+
return (
317+
SharedWeightInfo(
318+
name="word_embeddings",
319+
kind="tied_word_embeddings",
320+
canonical=canonical,
321+
aliases=(output,),
322+
),
323+
)
324+
325+
157326
def resolve_component_manifest(
158327
task: object,
159328
*,
@@ -198,4 +367,10 @@ def resolve_component_manifest(
198367
)
199368
for name in ordered_names
200369
)
201-
return ComponentManifest(descriptors)
370+
manifest = ComponentManifest(descriptors)
371+
return ComponentManifest(
372+
descriptors,
373+
shared_weights=_infer_shared_weights(hf_config, manifest)
374+
if hf_config is not None
375+
else (),
376+
)

‎src/mobius/_inspect.py‎

Lines changed: 24 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,11 +16,18 @@
1616

1717
from __future__ import annotations
1818

19-
__all__ = ["ComponentInfo", "inspect_components"]
19+
__all__ = [
20+
"ComponentInfo",
21+
"SharedWeightEndpoint",
22+
"SharedWeightInfo",
23+
"inspect_components",
24+
]
2025

2126
import dataclasses
2227
import logging
2328

29+
from mobius._component_manifest import SharedWeightEndpoint, SharedWeightInfo
30+
2431
logger = logging.getLogger(__name__)
2532

2633

@@ -45,11 +52,17 @@ class ComponentInfo:
4552
tuple. Empty when the component is the whole model or the layout is
4653
unknown. Tools such as Olive use these to optimize a submodule in
4754
place before exporting the full model.
55+
shared_weights: Cross-component shared parameters involving this
56+
component. The same immutable declaration is attached to every
57+
participating component. Automatic inference requires an
58+
unambiguous embedding/head pair with explicit source-path aliases;
59+
source paths alone do not establish weight sharing.
4860
"""
4961

5062
name: str
5163
role: str
5264
source_paths: tuple[str, ...] = ()
65+
shared_weights: tuple[SharedWeightInfo, ...] = ()
5366

5467

5568
def _resolve_task_model_type_and_config(
@@ -158,12 +171,22 @@ def inspect_components(
158171
model_type=model_type,
159172
hf_config=hf_config,
160173
)
174+
shared_weights = manifest.shared_weights
175+
shared_by_component = {
176+
name: tuple(
177+
shared_weight
178+
for shared_weight in shared_weights
179+
if any(endpoint.component == name for endpoint in shared_weight.endpoints)
180+
)
181+
for name in manifest
182+
}
161183

162184
components = [
163185
ComponentInfo(
164186
name=component.name,
165187
role=component.role,
166188
source_paths=component.source_paths,
189+
shared_weights=shared_by_component[component.name],
167190
)
168191
for component in manifest.values()
169192
]

0 commit comments

Comments
 (0)