88__all__ = [
99 "ComponentDescriptor" ,
1010 "ComponentManifest" ,
11+ "SharedWeightEndpoint" ,
12+ "SharedWeightInfo" ,
1113 "get_hf_component_sources" ,
1214 "resolve_component_manifest" ,
1315]
2022if 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 )
2582class 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+
157326def 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+ )
0 commit comments