Coverage for transformer_lens/model_bridge/component_setup.py: 88%
241 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-10-23 20:45 +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):
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 for parameter_name, parameter in module._parameters.items():
139 if parameter is not None and id(parameter) not in registered_parameter_ids:
140 missing_parameters.append((module_path, module, parameter_name, parameter))
141 registered_parameter_ids.add(id(parameter))
142 for buffer_name, buffer in module._buffers.items():
143 if buffer is not None and id(buffer) not in registered_buffer_ids:
144 missing_buffers.append((module_path, module, buffer_name, buffer))
145 registered_buffer_ids.add(id(buffer))
147 if not missing_parameters and not missing_buffers:
148 return
150 owner_by_path: dict[str, _ContainerStateOwner] = {"": _ContainerStateOwner(original_model)}
151 root_owner = owner_by_path[""]
152 original_modules = dict(original_model.named_modules())
154 def get_owner(module_path: str) -> _ContainerStateOwner:
155 current_path = ""
156 current_owner = root_owner
157 for path_part in module_path.split(".") if module_path else ():
158 child_path = f"{current_path}.{path_part}" if current_path else path_part
159 if child_path not in owner_by_path:
160 child_owner = _ContainerStateOwner(original_modules[child_path])
161 current_owner.add_module(path_part, child_owner)
162 owner_by_path[child_path] = child_owner
163 current_owner = owner_by_path[child_path]
164 current_path = child_path
165 return current_owner
167 for module_path, _, parameter_name, parameter in missing_parameters:
168 get_owner(module_path).register_parameter(parameter_name, parameter)
170 for module_path, module, buffer_name, buffer in missing_buffers:
171 get_owner(module_path).register_buffer(
172 buffer_name,
173 buffer,
174 persistent=buffer_name not in module._non_persistent_buffers_set,
175 )
177 bridge_module.add_module("_container_state_owners", root_owner)
180def _wire_symbolic_hooks(symbolic: GeneralizedComponent) -> None:
181 """Fire a placeholder's own hook_in/hook_out from its ``in``/``out`` subcomponents.
183 Applies to SymbolicBridge and to containerless bridges (``name=None``, e.g.
184 OPT/XGLM's fc-split MLPBridge): neither has a forward in the execution path,
185 so their own HookPoints would otherwise never fire. Permanent hooks (survive
186 reset_hooks) re-fire the subcomponent activation through the placeholder's
187 HookPoint, so ``blocks.{i}.mlp.hook_in/hook_out`` (and their compat aliases)
188 behave like every non-symbolic arch's. A hook that returns a modified tensor
189 on the placeholder point propagates: the mirror returns it into the
190 subcomponent's chain.
191 """
192 if getattr(symbolic, "_placeholder_hooks_wired", False): 192 ↛ 193line 192 didn't jump to line 193 because the condition on line 192 was never true
193 return
194 sub_in = symbolic.submodules.get("in")
195 sub_out = symbolic.submodules.get("out")
196 if sub_in is not None: 196 ↛ 198line 196 didn't jump to line 198 because the condition on line 196 was always true
197 sub_in.hook_in.add_hook(lambda tensor, hook: symbolic.hook_in(tensor), is_permanent=True)
198 if sub_out is not None: 198 ↛ 200line 198 didn't jump to line 200 because the condition on line 198 was always true
199 sub_out.hook_out.add_hook(lambda tensor, hook: symbolic.hook_out(tensor), is_permanent=True)
200 object.__setattr__(symbolic, "_placeholder_hooks_wired", True)
203def setup_submodules(
204 component: GeneralizedComponent,
205 architecture_adapter: ArchitectureAdapter,
206 original_model: RemoteModel,
207) -> None:
208 """Set up submodules for a bridge component using proper component setup.
210 Args:
211 component: The bridge component to set up submodules for
212 architecture_adapter: The architecture adapter
213 original_model: The original model to get components from
214 """
215 skipped_optional: list[str] = []
216 for module_name, submodule in component.submodules.items():
217 if submodule.is_list_item:
218 if submodule.name is None: 218 ↛ 219line 218 didn't jump to line 219 because the condition on line 218 was never true
219 raise ValueError(f"List item component {module_name} must have a name")
220 bridged_list = setup_blocks_bridge(submodule, architecture_adapter, original_model)
221 component.add_module(module_name, bridged_list)
222 replace_remote_component(bridged_list, submodule.name, original_model)
223 # Add to real_components mapping
224 component.real_components[module_name] = (submodule.name, list(bridged_list))
225 elif isinstance(submodule, SymbolicBridge): 225 ↛ 227line 225 didn't jump to line 227 because the condition on line 225 was never true
226 # SymbolicBridge: no real component; set up submodules via parent's model
227 setup_submodules(submodule, architecture_adapter, original_model)
229 # Add the symbolic bridge as a module (for structural access like blocks[i].mlp.in)
230 if module_name not in component._modules:
231 component.add_module(module_name, submodule)
233 # Add symbolic bridge's real_components to parent's mapping with prefixed keys
234 for sub_name, (sub_path, sub_comp) in submodule.real_components.items():
235 prefixed_key = f"{module_name}.{sub_name}"
236 component.real_components[prefixed_key] = (sub_path, sub_comp)
238 # The placeholder has no forward, so its own hook_in/hook_out would never
239 # fire — mirror them from the designated subcomponents (fc-split archs:
240 # mlp.hook_in = fc1's input, mlp.hook_out = fc2's output). Firing through
241 # the parent HookPoint keeps caching AND interventions working: the return
242 # value feeds back into the subcomponent's hook chain.
243 _wire_symbolic_hooks(submodule)
244 else:
245 # Set up original_component if not already set
246 if submodule.original_component is None:
247 if submodule.name is None:
248 original_subcomponent = original_model
249 else:
250 remote_path = submodule.name
251 is_optional = getattr(submodule, "optional", False)
252 # Fast path: first segment absent or None → skip
253 first_segment = remote_path.split(".")[0]
254 first_value = getattr(original_model, first_segment, None)
255 if is_optional and first_value is None:
256 logger.debug(
257 "Optional '%s' (path '%s') absent on %s",
258 module_name,
259 remote_path,
260 getattr(component, "name", "?"),
261 )
262 skipped_optional.append(module_name)
263 continue
264 # Full resolution — catches deeper path failures (e.g. stub self_attn missing q_proj)
265 try:
266 original_subcomponent = architecture_adapter.get_remote_component(
267 original_model, remote_path
268 )
269 except AttributeError:
270 if is_optional: 270 ↛ 271line 270 didn't jump to line 271 because the condition on line 270 was never true
271 logger.debug(
272 "Optional '%s' (path '%s') partially absent on %s",
273 module_name,
274 remote_path,
275 getattr(component, "name", "?"),
276 )
277 skipped_optional.append(module_name)
278 continue
279 raise
280 submodule.set_original_component(original_subcomponent)
281 setup_submodules(submodule, architecture_adapter, original_subcomponent)
282 if submodule.name is not None:
283 replace_remote_component(submodule, submodule.name, original_model)
285 # Add to _modules if not already present
286 if module_name not in component._modules:
287 component.add_module(module_name, submodule)
289 # Containerless executable bridges still own real child components.
290 # Promote them just like SymbolicBridge so processed weights reach
291 # projections that live directly on the parent block.
292 if not submodule.is_list_item and submodule.name is None:
293 for sub_name, (sub_path, sub_comp) in submodule.real_components.items():
294 prefixed_key = f"{module_name}.{sub_name}"
295 component.real_components[prefixed_key] = (sub_path, sub_comp)
296 # And like SymbolicBridge, nothing in the execution path calls
297 # their forward, so mirror hook_in/hook_out from in/out.
298 # Opt-in flag: executable containerless views (LLaDA's gated
299 # MLP) fire their own hooks and must not be double-fired.
300 if getattr(submodule, "mirror_placeholder_hooks", False):
301 _wire_symbolic_hooks(submodule)
302 elif not submodule.is_list_item: 302 ↛ 216line 302 didn't jump to line 216 because the condition on line 302 was always true
303 component.real_components[module_name] = (submodule.name, submodule)
305 # Clean up so architecture_adapter traversal won't find stale entries
306 for name in skipped_optional:
307 component.submodules.pop(name, None)
308 if skipped_optional:
309 _prune_hook_aliases_for_skipped(component, skipped_optional)
310 # Components whose submodules are optional only in some configurations
311 # validate here: which optionals were actually skipped is not knowable
312 # until every submodule has been resolved.
313 validate_after_setup = getattr(component, "validate_after_setup", None)
314 if callable(validate_after_setup):
315 validate_after_setup(skipped_optional)
318def _prune_hook_aliases_for_skipped(component: GeneralizedComponent, skipped: list[str]) -> None:
319 """Remove aliases targeting skipped optional submodules.
321 An overridden alias falls back to its class-level target when that target is
322 live on the current component.
323 """
324 aliases = getattr(component, "hook_aliases", None)
325 if not aliases:
326 return
327 skipped_set = set(skipped)
328 default_aliases = type(component).hook_aliases
329 # Deepcopied blocks still share the class-level dict until mutated.
330 if aliases is type(component).hook_aliases: 330 ↛ 334line 330 didn't jump to line 334 because the condition on line 330 was always true
331 aliases = dict(aliases)
332 component.hook_aliases = aliases
334 def _first_segment(path: str) -> str:
335 return path.split(".", 1)[0]
337 def _live_fallback(alias_name: str) -> str | list[str] | None:
338 fallback = default_aliases.get(alias_name)
339 if fallback is None:
340 return None
341 targets = fallback if isinstance(fallback, list) else [fallback]
342 live_targets = []
343 for target in targets:
344 if _first_segment(target) in skipped_set:
345 continue
346 current: nn.Module = component
347 for part in target.split("."):
348 modules = object.__getattribute__(current, "_modules")
349 next_component = modules.get(part)
350 if next_component is None: 350 ↛ 351line 350 didn't jump to line 351 because the condition on line 350 was never true
351 break
352 current = next_component
353 else:
354 live_targets.append(target)
355 if not live_targets:
356 return None
357 return live_targets if isinstance(fallback, list) else live_targets[0]
359 to_drop: list[str] = []
360 for alias_name, target in aliases.items():
361 if isinstance(target, list): 361 ↛ 362line 361 didn't jump to line 362 because the condition on line 361 was never true
362 kept = [t for t in target if _first_segment(t) not in skipped_set]
363 if not kept:
364 fallback = _live_fallback(alias_name)
365 if fallback is None:
366 to_drop.append(alias_name)
367 else:
368 aliases[alias_name] = fallback
369 elif len(kept) != len(target):
370 aliases[alias_name] = kept
371 elif _first_segment(target) in skipped_set:
372 fallback = _live_fallback(alias_name)
373 if fallback is None:
374 to_drop.append(alias_name)
375 else:
376 aliases[alias_name] = fallback
377 for alias_name in to_drop:
378 aliases.pop(alias_name, None)
381def setup_components(
382 components: dict[str, Any],
383 bridge_module: nn.Module,
384 architecture_adapter: ArchitectureAdapter,
385 original_model: RemoteModel,
386) -> None:
387 """Set up components on the bridge module.
389 Args:
390 components: Dictionary of component name to bridge component mappings
391 bridge_module: The bridge module to configure
392 architecture_adapter: The architecture adapter
393 original_model: The original model to get components from
394 """
395 for tl_path, bridge_component in components.items():
396 remote_path = bridge_component.name
397 if bridge_component.is_list_item:
398 bridged_list = setup_blocks_bridge(
399 bridge_component, architecture_adapter, original_model
400 )
401 bridge_module.add_module(tl_path, bridged_list)
402 replace_remote_component(bridged_list, remote_path, original_model)
403 # Add to bridge module's real_components if it has the attribute
404 if hasattr(bridge_module, "real_components"):
405 bridge_module.real_components[tl_path] = (remote_path, list(bridged_list)) # type: ignore[index, assignment, operator]
406 else:
407 original_component = architecture_adapter.get_remote_component(
408 original_model, remote_path
409 )
410 bridge_component.set_original_component(original_component)
411 setup_submodules(bridge_component, architecture_adapter, original_component)
412 bridge_module.add_module(tl_path, bridge_component)
413 replace_remote_component(bridge_component, remote_path, original_model)
414 # Add to bridge module's real_components if it has the attribute
415 if hasattr(bridge_module, "real_components"):
416 bridge_module.real_components[tl_path] = (remote_path, bridge_component) # type: ignore[index, assignment, operator]
419def setup_blocks_bridge(
420 blocks_template: Any, architecture_adapter: ArchitectureAdapter, original_model: RemoteModel
421) -> nn.ModuleList:
422 """Set up blocks bridge with proper ModuleList structure.
424 Args:
425 blocks_template: Template bridge component for blocks
426 architecture_adapter: The architecture adapter
427 original_model: The original model to get components from
429 Returns:
430 ModuleList of bridged block components
431 """
432 original_blocks = architecture_adapter.get_remote_component(
433 original_model, blocks_template.name
434 )
435 if not hasattr(original_blocks, "__iter__"): 435 ↛ 436line 435 didn't jump to line 436 because the condition on line 435 was never true
436 raise TypeError(f"Component {blocks_template.name} is not iterable")
437 bridged_blocks = nn.ModuleList()
438 iterable_blocks = cast(Any, original_blocks)
439 for i, original_block in enumerate(iterable_blocks):
440 block_bridge = copy.deepcopy(blocks_template)
441 block_bridge.name = f"{blocks_template.name}.{i}"
442 block_bridge.set_original_component(original_block)
443 setup_submodules(block_bridge, architecture_adapter, original_block)
444 if hasattr(block_bridge, "_wire_ln1_module"):
445 block_bridge._wire_ln1_module()
446 bridged_blocks.append(block_bridge)
447 replace_remote_component(bridged_blocks, blocks_template.name, original_model)
448 return bridged_blocks