Coverage for transformer_lens/model_bridge/generalized

1"""Base class for generalized transformer components."""

2from __future__ import annotations

4import contextlib

5import inspect

6import warnings

7from collections.abc import Callable

8from typing import Any, Dict, List, Optional, Union

10import torch

11import torch.nn as nn

12from accelerate.utils import align_module_device

14from transformer_lens.conversion_utils.conversion_steps.base_tensor_conversion import (

15 BaseTensorConversion,

16)

17from transformer_lens.hook_points import HookPoint

20class _ScaledHookGradient(torch.autograd.Function):

21 """Keep the raw forward value but differentiate through the scaled hook."""

23 generate_vmap_rule = True

25 @staticmethod

26 def forward(hooked: torch.Tensor, raw: torch.Tensor, scale: float) -> torch.Tensor:

27 return raw.clone()

29 @staticmethod

30 def setup_context(

31 ctx: Any, inputs: tuple[torch.Tensor, torch.Tensor, float], output: torch.Tensor

32 ) -> None:

33 ctx.scale = inputs[2]

35 @staticmethod

36 def backward(ctx: Any, grad: torch.Tensor) -> tuple[torch.Tensor, None, None]:

37 return grad / ctx.scale, None, None

39 @staticmethod

40 def jvp(

41 ctx: Any, grad_hooked: torch.Tensor, grad_raw: torch.Tensor, grad_scale: None

42 ) -> torch.Tensor:

43 return grad_hooked / ctx.scale

46def align_offloaded_subtree(module: nn.Module) -> contextlib.ExitStack:

47 """Materialize every Accelerate-offloaded descendant of ``module`` for the

48 caller's duration (an ``ExitStack`` of ``align_module_device`` contexts).

50 Accelerate attaches offload hooks at leaf level - whichever submodule

51 directly owns the Parameter (e.g. ``c_attn``, ``c_proj``) - not on

52 container modules like an attention block as a whole. ``align_module_device``

53 on a container alone is therefore a no-op even though its descendants are

54 offloaded. Walking every descendant and entering each one's

55 ``align_module_device`` (a cheap no-op for any module that has no hook of

56 its own) covers both a leaf ``original_component`` and a multi-level

57 container uniformly, without needing to know in advance which specific

58 descendant a given architecture adapter's code actually reads from.

59 """

60 stack = contextlib.ExitStack()

61 for submodule in module.modules():

62 stack.enter_context(align_module_device(submodule))

63 return stack

66class CloneOutputUnderGradMixin(nn.Module):

67 """Clone the forward output so HF's in-place mutation cannot corrupt it.

69 Under grad, autograd forbids in-place writes to backward-hook views; under

70 no_grad, cached hook_out tensors alias the storage HF then rewrites.

71 Mix in ahead of a bridge class: ``class X(CloneOutputUnderGradMixin, LinearBridge)``.

72 """

74 def forward(self, *args: Any, **kwargs: Any) -> Any:

75 out = super().forward(*args, **kwargs)

76 if isinstance(out, torch.Tensor): 76 ↛ 78line 76 didn't jump to line 78 because the condition on line 76 was always true

77 out = out.clone()

78 return out

81# Bumped whenever any component rebinds its hook_aliases (MoEBridge's dense/

82# sparse dispatch, OlmoHybrid's per-layer selection). Alias caches key on it:

83# a rebind can leave the hook REGISTRY unchanged while changing what the

84# aliases point at, so size-based cache keys cannot see it.

85_ALIAS_GENERATION = 0

88def alias_generation() -> int:

89 """Current global alias-rebind generation."""

90 return _ALIAS_GENERATION

93def _bump_alias_generation() -> None:

94 global _ALIAS_GENERATION

95 _ALIAS_GENERATION += 1

98class GeneralizedComponent(nn.Module):

99 """Base class for generalized transformer components.

100

101 This class provides a standardized interface for transformer components

102 and handles hook registration and execution.

103 """

104

105 is_list_item: bool = False

106 hook_out_is_single_residual_stream: bool = False

107 compatibility_mode: bool = False

108 disable_warnings: bool = False

109 hook_aliases: Dict[str, Union[str, List[str]]] = {}

110

111 # Projection-hook protocol between container bridges (MLPBridge) and the

112 # projection bridges they wrap (LinearBridge / Conv1DBridge): the wrapper

113 # records that its hook_out fired so the container only re-fires it as a

114 # bypass fallback, and the container suppresses the wrapper's next hook_in

115 # after pre-firing it itself. See MLPBridge.forward.

116 _fired_hook_out: bool = False

117 _suppress_next_hook_in: bool = False

118 property_aliases: Dict[str, str] = {}

119

120 def __init__(

121 self,

122 name: Optional[str],

123 config: Optional[Any] = None,

124 submodules: Optional[Dict[str, "GeneralizedComponent"]] = None,

125 conversion_rule: Optional[BaseTensorConversion] = None,

126 hook_alias_overrides: Optional[Dict[str, str]] = None,

127 optional: bool = False,

128 ):

129 """Initialize the generalized component.

130

131 Args:

132 name: The name of this component (None if component has no container in remote model)

133 config: Optional configuration object for the component

134 submodules: Dictionary of GeneralizedComponent submodules to register

135 conversion_rule: Optional conversion rule for this component's hooks

136 hook_alias_overrides: Optional dictionary to override default hook aliases.

137 For example, {"hook_attn_out": "ln1_post.hook_out"} will make hook_attn_out

138 point to ln1_post.hook_out instead of the default value in self.hook_aliases.

139 optional: If True, setup skips this subtree when absent (hybrid architectures).

140 """

141 super().__init__()

142 self.name = name

143 self.config = config

144 self.submodules = submodules or {}

145 self.conversion_rule = conversion_rule

146 self.optional = optional

147 self._hook_registry: Dict[str, HookPoint] = {}

148 self._hook_alias_registry: Dict[str, Union[str, List[str]]] = {}

149 self._property_alias_registry: Dict[str, str] = {}

150 self.hook_in = HookPoint()

151 self.hook_out = HookPoint()

152 # real_components maps TL keys to (remote_path, actual_instance) tuples

153 # For list components, actual_instance will be a list of component instances

154 self.real_components: Dict[str, tuple] = {}

155 if self.conversion_rule is not None:

156 self.hook_in.hook_conversion = self.conversion_rule

157 self.hook_out.hook_conversion = self.conversion_rule

158

159 # Copy class-level hook_aliases and apply any overrides

160 if hook_alias_overrides is not None:

161 self.hook_aliases = self.__class__.hook_aliases.copy()

162 self.hook_aliases.update(hook_alias_overrides)

163

164 def _hook_scaled_output(

165 self, output: torch.Tensor, hook_point: HookPoint, scale: float

166 ) -> torch.Tensor:

167 """Expose a residual contribution while leaving an outer HF scale intact."""

168 if scale == 1.0:

169 return hook_point(output)

170 if not hook_point.has_hooks(dir="both"):

171 return output

172 hooked = hook_point(output * scale)

173 # Recompute the reference because an in-place hook can mutate its input.

174 if not hook_point.has_hooks(dir="bwd") and torch.equal(hooked, output * scale):

175 if scale != 0.0 and torch.is_grad_enabled():

176 # Equal values can still carry a different intervention graph.

177 return _ScaledHookGradient.apply(hooked, output.detach(), scale)

178 return output

179 if scale == 0.0:

180 raise ValueError("Cannot edit or backward-hook a contribution with zero runtime scale")

181 return hooked / scale

182

183 def mark_pos_dim(self, pos_dim: int) -> None:

184 """Declare the position axis of every HookPoint this component owns.

185

186 Call it on a component whose activations do not carry position on dim 1

187 of the batched layout (e.g. a norm applied to ``[batch, heads, pos, d_head]``)

188 so ``pos_slice`` slices positions on all of its hooks, including ones a

189 subclass adds (``hook_normalized`` / ``hook_scale`` on a norm). Submodule

190 components are not touched; they declare their own layout.

191 """

192 for child in self.children():

193 if isinstance(child, HookPoint): 193 ↛ 192line 193 didn't jump to line 192 because the condition on line 193 was always true

194 child.pos_dim = pos_dim

195

196 def __call__(self, *args: Any, **kwargs: Any) -> Any:

197 """Run forward(), materializing the wrapped component's params first if offloaded.

198

199 Bridge components read the wrapped module's raw parameters directly

200 (self.weight / self.original_component.bias / etc. via __getattr__ or

201 direct attribute access) rather than exclusively calling the wrapped

202 module's own forward(). Under an Accelerate CPU/disk device_map,

203 Accelerate only swaps a meta placeholder for the real, materialized

204 tensor around the wrapped module's OWN forward() call (its pre/post

205 forward hooks) - a raw attribute read outside that window silently

206 sees the meta placeholder instead, with no error. align_module_device

207 wraps this call in the same pre/post-forward hook Accelerate would

208 have run, so every read during this call - whichever way the code

209 reaches it - gets real data. It's a no-op (bare yield) when the

210 wrapped module has no offload hook, which covers the common case of

211 no device_map, a single device, or a multi-GPU split with no CPU/disk

212 offload involved.

213

214 Deliberately just original_component, not align_offloaded_subtree's

215 whole-subtree walk: Accelerate hooks the leaf modules that actually own

216 Parameters (e.g. NormalizationBridge/LinearBridge wrap one directly), so

217 materializing exactly that leaf for exactly its own call keeps the same

218 one-leaf-at-a-time memory footprint Accelerate's native per-module hooks

219 would give a plain (non-bridge) forward pass. A component whose

220 original_component is a container of several separately-hooked leaves

221 (e.g. an attention module wrapping distinct q/k/v/o projections) doesn't

222 need this to also materialize here - each of ITS own sub-bridges calls

223 into its own leaf the same way. align_offloaded_subtree is for the

224 narrower case of code that reads a specific descendant's raw params

225 directly during setup, without going through that descendant's own

226 bridge __call__ at all (see JointQKVAttentionBridge.set_original_component).

227 """

228 original_component = self._modules.get("_original_component")

229 if original_component is None:

230 return super().__call__(*args, **kwargs)

231 with align_module_device(original_component):

232 return super().__call__(*args, **kwargs)

233

234 def _register_hook(self, name: str, hook: HookPoint) -> None:

235 """Register a hook in the component's hook registry."""

236 hook.name = name

237 self._hook_registry[name] = hook

238

239 def _register_aliases(self) -> None:

240 """Register aliases from class-level dictionaries.

241

242 Called unconditionally at bridge init (see bridge.py); compatibility mode

243 additionally re-registers after weight processing.

244 It creates actual Python attributes/properties that directly reference the target objects.

245

246 Note: Re-registration expects to run after

247 weight processing is complete to ensure property aliases point to processed weights.

248 """

249 if self.hook_aliases:

250 self._hook_alias_registry.update(self.hook_aliases)

251 if self.property_aliases:

252 self._property_alias_registry.update(self.property_aliases)

253 for alias_name, target_path in self._hook_alias_registry.items():

254 resolved = False

255 if isinstance(target_path, list): 255 ↛ 256line 255 didn't jump to line 256 because the condition on line 255 was never true

256 for single_target in target_path:

257 try:

258 target_obj = self

259 for part in single_target.split("."):

260 target_obj = getattr(target_obj, part)

261 object.__setattr__(self, alias_name, target_obj)

262 resolved = True

263 break

264 except AttributeError:

265 continue

266 else:

267 try:

268 target_obj = self

269 for part in target_path.split("."):

270 target_obj = getattr(target_obj, part)

271 object.__setattr__(self, alias_name, target_obj)

272 resolved = True

273 except AttributeError:

274 pass

275 if not resolved:

276 # Surface drops instead of silently swallowing — some aliases are

277 # legitimately conditional on optional submodules, but an author

278 # needs to see which ones dropped at bridge-init.

279 warnings.warn(

280 f"Hook alias '{alias_name}' -> '{target_path}' on "

281 f"{type(self).__name__}(name={getattr(self, 'name', None)!r}) "

282 f"did not resolve; this hook will not be accessible.",

283 stacklevel=2,

284 )

285 for alias_name, target_path in self._property_alias_registry.items():

286 try:

287 target_obj = self

288 for part in target_path.split("."):

289 target_obj = getattr(target_obj, part)

290 object.__setattr__(self, alias_name, target_obj)

291 except AttributeError:

292 pass

293

294 def get_hooks(self) -> Dict[str, HookPoint]:

295 """Get all hooks registered in this component."""

296 hooks = self._hook_registry.copy()

297 if self.compatibility_mode and self._hook_alias_registry:

298 for alias_name in self._hook_alias_registry.keys():

299 if hasattr(self, alias_name): 299 ↛ 298line 299 didn't jump to line 298 because the condition on line 299 was always true

300 target_hook = getattr(self, alias_name)

301 if isinstance(target_hook, HookPoint): 301 ↛ 298line 301 didn't jump to line 298 because the condition on line 301 was always true

302 hooks[alias_name] = target_hook

303 return hooks

304

305 def _is_getattr_called_internally(self) -> bool:

306 """This function checks if the __getattr__ method was being called internally

307 (e.g by the setup process or run_with_cache).

308 """

309 for frame_info in inspect.stack():

310 if "setup_components" in frame_info.function or "run_with_cache" in frame_info.function:

311 return True

312 return False

313

314 def set_original_component(self, original_component: nn.Module) -> None:

315 """Set the original component that this bridge wraps.

316

317 Args:

318 original_component: The original transformer component to wrap

319 """

320 self.add_module("_original_component", original_component)

321 # An opaque wrapper (created with config=None) shadows the wrapped

322 # module's own config. HF forwards sometimes read a submodule's config

323 # directly (e.g. Qwen2Audio's forward does create_bidirectional_mask(

324 # config=self.audio_tower.config)), so inherit the real config to avoid

325 # exposing None. Components given an explicit config keep it.

326 if self.config is None:

327 self.config = getattr(original_component, "config", None)

328

329 @property

330 def original_component(self) -> Optional[nn.Module]:

331 """Get the original component."""

332 return self._modules.get("_original_component", None)

333

334 def add_hook(self, hook_fn: Callable[..., torch.Tensor], hook_name: str = "output") -> None:

335 """Add a hook function (HookedTransformer-compatible interface).

336

337 Args:

338 hook_fn: Function to call at this hook point

339 hook_name: Name of the hook point (defaults to "output")

340 """

341 if hook_name == "output":

342 self.hook_out.add_hook(hook_fn)

343 elif hook_name == "input":

344 self.hook_in.add_hook(hook_fn)

345 else:

346 raise ValueError(

347 f"Hook name '{hook_name}' not supported. Supported names are 'output' and 'input'."

348 )

349

350 def remove_hooks(self, hook_name: str | None = None) -> None:

351 """Remove hooks (HookedTransformer-compatible interface).

352

353 Args:

354 hook_name: Name of the hook point to remove. If None, removes all hooks.

355 """

356 if hook_name is None:

357 self.hook_in.remove_hooks(dir="both")

358 self.hook_out.remove_hooks(dir="both")

359 elif hook_name == "output":

360 self.hook_out.remove_hooks(dir="both")

361 elif hook_name == "input":

362 self.hook_in.remove_hooks(dir="both")

363 else:

364 raise ValueError(

365 f"Hook name '{hook_name}' not supported. Supported names are 'output' and 'input'."

366 )

367

368 def set_processed_weights(

369 self, weights: Dict[str, torch.Tensor], verbose: bool = False

370 ) -> None:

371 """Set the processed weights for use in compatibility mode.

372

373 This method stores processed weights as attributes on the component so they can be

374 used directly in the forward pass without modifying the original component.

375

376 Components should override this method to handle their specific weight structure.

377 The weights dict contains keys like "weight", "bias", "W_in", "W_out", etc.

378

379 If this component has submodules, this method will automatically distribute the

380 weights to those subcomponents using ProcessWeights.distribute_weights_to_components.

381

382 Args:

383 weights: Dictionary of processed weight tensors

384 verbose: If True, print detailed information about weight setting

385 """

386 if verbose: 386 ↛ 387line 386 didn't jump to line 387 because the condition on line 386 was never true

387 print(

388 f"\n set_processed_weights: {self.__class__.__name__} (name={getattr(self, 'name', 'unknown')})"

389 )

390 print(f" Received {len(weights)} weight keys")

391

392 # First, handle single-part keys (keys without ".") by setting them as parameters

393 # on the original component

394 if self.original_component is not None:

395 for key, weight_tensor in weights.items():

396 # Only process keys without "." (single-part keys)

397 if "." not in key:

398 # Try to set the parameter on the original component

399 if hasattr(self.original_component, key):

400 param = getattr(self.original_component, key)

401 if param is not None and isinstance(param, torch.nn.Parameter):

402 if param.shape != weight_tensor.shape: 402 ↛ 403line 402 didn't jump to line 403 because the condition on line 402 was never true

403 raise ValueError(

404 f"Shape mismatch when setting weight '{key}' in {type(self.original_component).__name__}: "

405 f"existing param shape {param.shape} != new tensor shape {weight_tensor.shape}"

406 )

407 if verbose: 407 ↛ 408line 407 didn't jump to line 408 because the condition on line 407 was never true

408 print(f" Setting weight: {key} (shape: {weight_tensor.shape})")

409 # break tying by creating a new param

410 new_param = nn.Parameter(weight_tensor)

411 setattr(self.original_component, key, new_param)

412 elif param is None:

413 # Parameter exists but is None (e.g., bias=False in nn.Linear)

414 # Create a new parameter from the weight tensor

415 if verbose: 415 ↛ 416line 415 didn't jump to line 416 because the condition on line 415 was never true

416 print(

417 f" Creating weight: {key} (shape: {weight_tensor.shape}) - was None"

418 )

419 new_param = nn.Parameter(weight_tensor)

420 setattr(self.original_component, key, new_param)

421

422 # If this component has submodules, distribute weights to them

423 if self.real_components:

424 from transformer_lens.weight_processing import ProcessWeights

425

426 if verbose: 426 ↛ 427line 426 didn't jump to line 427 because the condition on line 426 was never true

427 print(f" Has {len(self.real_components)} subcomponents, distributing weights...")

428

429 ProcessWeights.distribute_weights_to_components(

430 state_dict=weights,

431 component_mapping=self.real_components,

432 verbose=verbose,

433 )

434

435 def forward(self, *args: Any, **kwargs: Any) -> Any:

436 """Generic forward pass for bridge components with input/output hooks."""

437 original_component = self._modules.get("_original_component", None)

438 if original_component is None: 438 ↛ 439line 438 didn't jump to line 439 because the condition on line 438 was never true

439 raise RuntimeError(

440 f"Original component not set for {self.name}. Call set_original_component() first."

441 )

442 input_arg_names = [

443 "input",

444 "hidden_states",

445 "input_ids",

446 "query_input",

447 "x",

448 "inputs_embeds",

449 ]

450 input_found = False

451 for name in input_arg_names:

452 if name in kwargs:

453 hooked = self.hook_in(kwargs[name])

454 kwargs[name] = hooked

455 input_found = True

456 break

457 if not input_found and len(args) > 0 and isinstance(args[0], torch.Tensor):

458 hooked_input = self.hook_in(args[0])

459 args = (hooked_input,) + args[1:]

460 input_found = True

461 output = original_component(*args, **kwargs)

462 if isinstance(output, tuple):

463 hooked_first = self.hook_out(output[0])

464 output = (hooked_first,) + output[1:]

465 elif not isinstance(output, torch.Tensor) and isinstance(

466 getattr(output, "last_hidden_state", None), torch.Tensor

467 ):

468 # ModelOutput-returning components (e.g. vision/audio towers).

469 output.last_hidden_state = self.hook_out(output.last_hidden_state)

470 else:

471 output = self.hook_out(output)

472 return output

473

474 def __getattr__(self, name: str) -> Any:

475 modules = object.__getattribute__(self, "__dict__").get("_modules")

476 if modules is not None and name in modules:

477 return modules[name]

478 if name == "original_component": 478 ↛ 479line 478 didn't jump to line 479 because the condition on line 478 was never true

479 raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'")

480 submodules = object.__getattribute__(self, "__dict__").get("submodules")

481 if submodules is not None and name in submodules:

482 # Don't return submodule here - it should be accessed via _modules after add_module()

483 # Raising AttributeError allows PyTorch's add_module() to work correctly

484 raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'")

485 if modules is not None: 485 ↛ 499line 485 didn't jump to line 499 because the condition on line 485 was always true

486 original_component = modules.get("_original_component")

487 if original_component is not None:

488 try:

489 if "." in name: 489 ↛ 490line 489 didn't jump to line 490 because the condition on line 489 was never true

490 name_split = name.split(".")

491 current = getattr(original_component, name_split[0])

492 for part in name_split[1:]:

493 current = getattr(current, part)

494 return current

495 else:

496 return getattr(original_component, name)

497 except AttributeError:

498 pass

499 raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'")

500

501 def __setattr__(self, name: str, value: Any) -> None:

502 """Set attribute, with passthrough to original component for compatibility."""

503 if isinstance(value, HookPoint):

504 self._register_hook(name, value)

505 super().__setattr__(name, value)

506 return

507 if name == "hook_aliases":

508 # Any alias rebind invalidates alias caches downstream.

509 _bump_alias_generation()

510 if name.startswith("_") or name in [

511 "name",

512 "config",

513 "submodules",

514 "conversion_rule",

515 "compatibility_mode",

516 "disable_warnings",

517 "optional",

518 # Components rebind these per layer at bind time (MoEBridge's

519 # dense/sparse dispatch). Without the carve-out the assignment is

520 # forwarded to the wrapped HF module whenever it happens to expose

521 # the attribute — the rebind then silently vanishes.

522 "hook_aliases",

523 "property_aliases",

524 # train()/eval() set self.training; redirecting it to the original

525 # component leaves the wrapper stuck in training mode (dropout at

526 # inference). Recursion still reaches the original via _modules.

527 "training",

528 ]:

529 super().__setattr__(name, value)

530 return

531 class_attr = getattr(type(self), name, None)

532 if class_attr is not None and isinstance(class_attr, property):

533 if class_attr.fset is not None:

534 super().__setattr__(name, value)

535 return

536 if hasattr(self, "_modules") and "_original_component" in self._modules:

537 original_component = self._modules["_original_component"]

538 if hasattr(original_component, name):

539 try:

540 setattr(original_component, name, value)

541 return

542 except AttributeError:

543 pass

544 super().__setattr__(name, value)

Coverage for transformer_lens/model_bridge/generalized_components/base.py: 87%

264 statements