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