Coverage for transformer_lens/model_bridge/component_setup.py: 88%
229 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-09-01 16:23 +0000
1from __future__ import annotations
3"Component setup utilities for creating and configuring bridged components."
4import copy
5import logging
6from typing import TYPE_CHECKING, Any, cast
8logger = logging.getLogger(__name__)
10import torch.nn as nn
12from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter
13from transformer_lens.model_bridge.generalized_components.base import (
14 GeneralizedComponent,
15)
16from transformer_lens.model_bridge.generalized_components.symbolic import SymbolicBridge
17from transformer_lens.model_bridge.types import RemoteModel
19if TYPE_CHECKING:
20 pass
23class _ContainerStateOwner(nn.Module):
24 """Registered view of state owned directly by an unwrapped container."""
26 def __init__(self, original_container: nn.Module) -> None:
27 super().__init__()
28 self.__dict__["_original_container"] = original_container
30 def _sync_original_container(self) -> None:
31 original_container = self.__dict__["_original_container"]
32 original_container._parameters.update(self._parameters)
33 original_container._buffers.update(self._buffers)
35 def _refresh_from_original_container(self) -> None:
36 original_container = self.__dict__["_original_container"]
37 for name in self._parameters:
38 self._parameters[name] = original_container._parameters[name]
39 for name in self._buffers:
40 self._buffers[name] = original_container._buffers[name]
42 def _apply(self, fn: Any, recurse: bool = True) -> "_ContainerStateOwner":
43 self._refresh_from_original_container()
44 super()._apply(fn, recurse=recurse)
45 self._sync_original_container()
46 return self
48 def _load_from_state_dict(
49 self,
50 state_dict: dict[str, Any],
51 prefix: str,
52 local_metadata: dict[str, Any],
53 strict: bool,
54 missing_keys: list[str],
55 unexpected_keys: list[str],
56 error_msgs: list[str],
57 ) -> None:
58 super()._load_from_state_dict(
59 state_dict,
60 prefix,
61 local_metadata,
62 strict,
63 missing_keys,
64 unexpected_keys,
65 error_msgs,
66 )
67 self._sync_original_container()
70def refresh_container_state_owners(bridge_module: nn.Module) -> None:
71 """Refresh registered container state from the original model tree."""
72 root_owner = bridge_module._modules.get("_container_state_owners")
73 if not isinstance(root_owner, _ContainerStateOwner): 73 ↛ 74line 73 didn't jump to line 74 because the condition on line 73 was never true
74 return
75 for owner in root_owner.modules():
76 if isinstance(owner, _ContainerStateOwner): 76 ↛ 75line 76 didn't jump to line 75 because the condition on line 76 was always true
77 owner._refresh_from_original_container()
80def replace_remote_component(
81 replacement_component: nn.Module, remote_path: str, remote_model: RemoteModel
82) -> None:
83 """Replace a component in a remote model.
85 Args:
86 replacement_component: The new component to install
87 remote_path: Path to the component in the remote model
88 remote_model: The remote model to modify
89 """
90 path_parts = remote_path.split(".")
91 current = remote_model
92 for part in path_parts[:-1]:
93 if hasattr(current, part):
94 current = getattr(current, part)
95 else:
96 raise ValueError(f"Path {remote_path} not found in model")
97 target_attr = path_parts[-1]
98 if hasattr(current, target_attr):
99 setattr(current, target_attr, replacement_component)
100 else:
101 raise ValueError(f"Attribute {target_attr} not found in {current}")
104def set_original_components(
105 bridge_module: nn.Module, architecture_adapter: ArchitectureAdapter, original_model: RemoteModel
106) -> None:
107 """Set original components on the pre-created bridge components.
109 Args:
110 bridge_module: The bridge module to configure
111 architecture_adapter: The architecture adapter
112 original_model: The original model to get components from
113 """
114 component_mapping = architecture_adapter.get_component_mapping()
115 setup_components(component_mapping, bridge_module, architecture_adapter, original_model)
116 if isinstance(original_model, nn.Module): 116 ↛ exitline 116 didn't return from function 'set_original_components' because the condition on line 116 was always true
117 _register_unowned_container_state(bridge_module, original_model)
120def _register_unowned_container_state(bridge_module: nn.Module, original_model: nn.Module) -> None:
121 """Make direct container parameters and buffers reachable through the Bridge tree."""
122 registered_parameter_ids = {
123 id(parameter)
124 for module in bridge_module.modules()
125 for parameter in module._parameters.values()
126 if parameter is not None
127 }
128 registered_buffer_ids = {
129 id(buffer)
130 for module in bridge_module.modules()
131 for buffer in module._buffers.values()
132 if buffer is not None
133 }
134 missing_parameters: list[tuple[str, nn.Module, str, nn.Parameter]] = []
135 missing_buffers: list[tuple[str, nn.Module, str, Any]] = []
137 for module_path, module in original_model.named_modules():
138 if any(child is not None for child in module._modules.values()):
139 for parameter_name, parameter in module._parameters.items():
140 if parameter is not None and id(parameter) not in registered_parameter_ids:
141 missing_parameters.append((module_path, module, parameter_name, parameter))
142 registered_parameter_ids.add(id(parameter))
143 for buffer_name, buffer in module._buffers.items():
144 if buffer is not None and id(buffer) not in registered_buffer_ids:
145 missing_buffers.append((module_path, module, buffer_name, buffer))
146 registered_buffer_ids.add(id(buffer))
148 if not missing_parameters and not missing_buffers:
149 return
151 owner_by_path: dict[str, _ContainerStateOwner] = {"": _ContainerStateOwner(original_model)}
152 root_owner = owner_by_path[""]
153 original_modules = dict(original_model.named_modules())
155 def get_owner(module_path: str) -> _ContainerStateOwner:
156 current_path = ""
157 current_owner = root_owner
158 for path_part in module_path.split(".") if module_path else ():
159 child_path = f"{current_path}.{path_part}" if current_path else path_part
160 if child_path not in owner_by_path:
161 child_owner = _ContainerStateOwner(original_modules[child_path])
162 current_owner.add_module(path_part, child_owner)
163 owner_by_path[child_path] = child_owner
164 current_owner = owner_by_path[child_path]
165 current_path = child_path
166 return current_owner
168 for module_path, _, parameter_name, parameter in missing_parameters:
169 get_owner(module_path).register_parameter(parameter_name, parameter)
171 for module_path, module, buffer_name, buffer in missing_buffers:
172 get_owner(module_path).register_buffer(
173 buffer_name,
174 buffer,
175 persistent=buffer_name not in module._non_persistent_buffers_set,
176 )
178 bridge_module.add_module("_container_state_owners", root_owner)
181def setup_submodules(
182 component: GeneralizedComponent,
183 architecture_adapter: ArchitectureAdapter,
184 original_model: RemoteModel,
185) -> None:
186 """Set up submodules for a bridge component using proper component setup.
188 Args:
189 component: The bridge component to set up submodules for
190 architecture_adapter: The architecture adapter
191 original_model: The original model to get components from
192 """
193 skipped_optional: list[str] = []
194 for module_name, submodule in component.submodules.items():
195 if submodule.is_list_item:
196 if submodule.name is None: 196 ↛ 197line 196 didn't jump to line 197 because the condition on line 196 was never true
197 raise ValueError(f"List item component {module_name} must have a name")
198 bridged_list = setup_blocks_bridge(submodule, architecture_adapter, original_model)
199 component.add_module(module_name, bridged_list)
200 replace_remote_component(bridged_list, submodule.name, original_model)
201 # Add to real_components mapping
202 component.real_components[module_name] = (submodule.name, list(bridged_list))
203 elif isinstance(submodule, SymbolicBridge): 203 ↛ 205line 203 didn't jump to line 205 because the condition on line 203 was never true
204 # SymbolicBridge: no real component; set up submodules via parent's model
205 setup_submodules(submodule, architecture_adapter, original_model)
207 # Add the symbolic bridge as a module (for structural access like blocks[i].mlp.in)
208 if module_name not in component._modules:
209 component.add_module(module_name, submodule)
211 # Add symbolic bridge's real_components to parent's mapping with prefixed keys
212 for sub_name, (sub_path, sub_comp) in submodule.real_components.items():
213 prefixed_key = f"{module_name}.{sub_name}"
214 component.real_components[prefixed_key] = (sub_path, sub_comp)
215 else:
216 # Set up original_component if not already set
217 if submodule.original_component is None:
218 if submodule.name is None:
219 original_subcomponent = original_model
220 else:
221 remote_path = submodule.name
222 is_optional = getattr(submodule, "optional", False)
223 # Fast path: first segment absent or None → skip
224 first_segment = remote_path.split(".")[0]
225 first_value = getattr(original_model, first_segment, None)
226 if is_optional and first_value is None:
227 logger.debug(
228 "Optional '%s' (path '%s') absent on %s",
229 module_name,
230 remote_path,
231 getattr(component, "name", "?"),
232 )
233 skipped_optional.append(module_name)
234 continue
235 # Full resolution — catches deeper path failures (e.g. stub self_attn missing q_proj)
236 try:
237 original_subcomponent = architecture_adapter.get_remote_component(
238 original_model, remote_path
239 )
240 except AttributeError:
241 if is_optional: 241 ↛ 242line 241 didn't jump to line 242 because the condition on line 241 was never true
242 logger.debug(
243 "Optional '%s' (path '%s') partially absent on %s",
244 module_name,
245 remote_path,
246 getattr(component, "name", "?"),
247 )
248 skipped_optional.append(module_name)
249 continue
250 raise
251 submodule.set_original_component(original_subcomponent)
252 setup_submodules(submodule, architecture_adapter, original_subcomponent)
253 if submodule.name is not None:
254 replace_remote_component(submodule, submodule.name, original_model)
256 # Add to _modules if not already present
257 if module_name not in component._modules:
258 component.add_module(module_name, submodule)
260 # Containerless executable bridges still own real child components.
261 # Promote them just like SymbolicBridge so processed weights reach
262 # projections that live directly on the parent block.
263 if not submodule.is_list_item and submodule.name is None:
264 for sub_name, (sub_path, sub_comp) in submodule.real_components.items():
265 prefixed_key = f"{module_name}.{sub_name}"
266 component.real_components[prefixed_key] = (sub_path, sub_comp)
267 elif not submodule.is_list_item: 267 ↛ 194line 267 didn't jump to line 194 because the condition on line 267 was always true
268 component.real_components[module_name] = (submodule.name, submodule)
270 # Clean up so architecture_adapter traversal won't find stale entries
271 for name in skipped_optional:
272 component.submodules.pop(name, None)
273 if skipped_optional:
274 _prune_hook_aliases_for_skipped(component, skipped_optional)
275 # Components whose submodules are optional only in some configurations
276 # validate here: which optionals were actually skipped is not knowable
277 # until every submodule has been resolved.
278 validate_after_setup = getattr(component, "validate_after_setup", None)
279 if callable(validate_after_setup):
280 validate_after_setup(skipped_optional)
283def _prune_hook_aliases_for_skipped(component: GeneralizedComponent, skipped: list[str]) -> None:
284 """Remove aliases targeting skipped optional submodules.
286 An overridden alias falls back to its class-level target when that target is
287 live on the current component.
288 """
289 aliases = getattr(component, "hook_aliases", None)
290 if not aliases:
291 return
292 skipped_set = set(skipped)
293 default_aliases = type(component).hook_aliases
294 # Deepcopied blocks still share the class-level dict until mutated.
295 if aliases is type(component).hook_aliases:
296 aliases = dict(aliases)
297 component.hook_aliases = aliases
299 def _first_segment(path: str) -> str:
300 return path.split(".", 1)[0]
302 def _live_fallback(alias_name: str) -> str | list[str] | None:
303 fallback = default_aliases.get(alias_name)
304 if fallback is None:
305 return None
306 targets = fallback if isinstance(fallback, list) else [fallback]
307 live_targets = []
308 for target in targets:
309 if _first_segment(target) in skipped_set:
310 continue
311 current: nn.Module = component
312 for part in target.split("."):
313 modules = object.__getattribute__(current, "_modules")
314 next_component = modules.get(part)
315 if next_component is None: 315 ↛ 316line 315 didn't jump to line 316 because the condition on line 315 was never true
316 break
317 current = next_component
318 else:
319 live_targets.append(target)
320 if not live_targets:
321 return None
322 return live_targets if isinstance(fallback, list) else live_targets[0]
324 to_drop: list[str] = []
325 for alias_name, target in aliases.items():
326 if isinstance(target, list): 326 ↛ 327line 326 didn't jump to line 327 because the condition on line 326 was never true
327 kept = [t for t in target if _first_segment(t) not in skipped_set]
328 if not kept:
329 fallback = _live_fallback(alias_name)
330 if fallback is None:
331 to_drop.append(alias_name)
332 else:
333 aliases[alias_name] = fallback
334 elif len(kept) != len(target):
335 aliases[alias_name] = kept
336 elif _first_segment(target) in skipped_set:
337 fallback = _live_fallback(alias_name)
338 if fallback is None:
339 to_drop.append(alias_name)
340 else:
341 aliases[alias_name] = fallback
342 for alias_name in to_drop:
343 aliases.pop(alias_name, None)
346def setup_components(
347 components: dict[str, Any],
348 bridge_module: nn.Module,
349 architecture_adapter: ArchitectureAdapter,
350 original_model: RemoteModel,
351) -> None:
352 """Set up components on the bridge module.
354 Args:
355 components: Dictionary of component name to bridge component mappings
356 bridge_module: The bridge module to configure
357 architecture_adapter: The architecture adapter
358 original_model: The original model to get components from
359 """
360 for tl_path, bridge_component in components.items():
361 remote_path = bridge_component.name
362 if bridge_component.is_list_item:
363 bridged_list = setup_blocks_bridge(
364 bridge_component, architecture_adapter, original_model
365 )
366 bridge_module.add_module(tl_path, bridged_list)
367 replace_remote_component(bridged_list, remote_path, original_model)
368 # Add to bridge module's real_components if it has the attribute
369 if hasattr(bridge_module, "real_components"):
370 bridge_module.real_components[tl_path] = (remote_path, list(bridged_list)) # type: ignore[index, assignment, operator]
371 else:
372 original_component = architecture_adapter.get_remote_component(
373 original_model, remote_path
374 )
375 bridge_component.set_original_component(original_component)
376 setup_submodules(bridge_component, architecture_adapter, original_component)
377 bridge_module.add_module(tl_path, bridge_component)
378 replace_remote_component(bridge_component, remote_path, original_model)
379 # Add to bridge module's real_components if it has the attribute
380 if hasattr(bridge_module, "real_components"):
381 bridge_module.real_components[tl_path] = (remote_path, bridge_component) # type: ignore[index, assignment, operator]
384def setup_blocks_bridge(
385 blocks_template: Any, architecture_adapter: ArchitectureAdapter, original_model: RemoteModel
386) -> nn.ModuleList:
387 """Set up blocks bridge with proper ModuleList structure.
389 Args:
390 blocks_template: Template bridge component for blocks
391 architecture_adapter: The architecture adapter
392 original_model: The original model to get components from
394 Returns:
395 ModuleList of bridged block components
396 """
397 original_blocks = architecture_adapter.get_remote_component(
398 original_model, blocks_template.name
399 )
400 if not hasattr(original_blocks, "__iter__"): 400 ↛ 401line 400 didn't jump to line 401 because the condition on line 400 was never true
401 raise TypeError(f"Component {blocks_template.name} is not iterable")
402 bridged_blocks = nn.ModuleList()
403 iterable_blocks = cast(Any, original_blocks)
404 for i, original_block in enumerate(iterable_blocks):
405 block_bridge = copy.deepcopy(blocks_template)
406 block_bridge.name = f"{blocks_template.name}.{i}"
407 block_bridge.set_original_component(original_block)
408 setup_submodules(block_bridge, architecture_adapter, original_block)
409 if hasattr(block_bridge, "_wire_ln1_module"):
410 block_bridge._wire_ln1_module()
411 bridged_blocks.append(block_bridge)
412 replace_remote_component(bridged_blocks, blocks_template.name, original_model)
413 return bridged_blocks