Coverage for transformer_lens/model_bridge/supported_architectures/pretrain.py: 88%
108 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-09-21 19:27 +0000
1"""Architecture adapter for a lightweight decoder-only pretraining model.
3Maps a decoder-only transformer using RoPE, RMSNorm, gated SwiGLU MLPs, and
4optional sparse mixture-of-experts feed-forward layers into
5TransformerBridge, by wrapping the source module and delegating to its own
6`forward` rather than translating parameters into a second implementation.
8Usage: `build_pretrain_bridge(model, cfg)` -- the public entry point.
9`PretrainModelContainer` and direct `build_bridge_from_module` use are
10internal/advanced details (see `PretrainModelContainer`'s docstring).
12Scope: maps a live module into TransformerBridge. Does not load
13checkpoints, merge tensor-parallel shards, or depend on a training
14framework.
16Required module protocol -- "lightweight decoder-only pretraining models"
17describes intent, not a generality guarantee. The wrapped model must
18expose:
20 model.embed (embedding lookup)
21 model.blocks[i].norm1 (pre-attention norm)
22 model.blocks[i].attn (called as attn(x, ...))
23 model.blocks[i].norm2 (pre-MLP norm)
24 model.blocks[i].mlp (gate/up/down, or router/experts)
25 model.norm_f (final norm)
26 model.lm_head (unembedding)
28`gate`/`up`/`down` and `router`/`experts` name the supported protocol.
29`DenseOrMoEFeedForwardBridge` checks these structurally -- attribute
30presence plus basic type (each is a module, `experts` is a registered
31module collection) -- and raises clearly on a mismatch, but that is
32structural validation only: it does not and cannot validate that a
33module satisfying the shape actually implements matching forward
34semantics. Blocks must take more than the bare hidden state (this target
35passes `cos`/`sin`) -- see `PretrainModelContainer`.
36"""
37from __future__ import annotations
39from typing import Any
41import torch
43from transformer_lens.config import TransformerBridgeConfig
44from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter
45from transformer_lens.model_bridge.generalized_components import (
46 AttentionBridge,
47 DelegatedAttentionBlockBridge,
48 EmbeddingBridge,
49 GatedMLPBridge,
50 LinearBridge,
51 MoEBridge,
52 RMSNormalizationBridge,
53 UnembeddingBridge,
54)
55from transformer_lens.model_bridge.generalized_components.base import (
56 GeneralizedComponent,
57)
58from transformer_lens.model_bridge.transformer_bridge import TransformerBridge
60ARCHITECTURE_NAME = "TransformerLensPretrain"
62# Reserved bridge kwargs removed before forwarding to the wrapped model.
63# Only known bridge-compatibility kwargs are stripped so genuine caller
64# mistakes (e.g. `target=` vs `targets=`) still raise naturally, and no
65# signature introspection is needed to support an arbitrary forward.
66_BRIDGE_COMPAT_KWARGS = frozenset({"output_attentions"})
69class DenseOrMoEFeedForwardBridge(GeneralizedComponent):
70 """Wraps a dense SwiGLU MLP or sparse MoE layer behind one interface.
71 Dispatch is by structural inspection (`router`/`experts` vs
72 `gate`/`up`/`down`), not config, so dense/MoE/mixed architectures all
73 share the same component mapping.
74 """
76 def __init__(self, name: str, config: Any):
77 super().__init__(name, config=config, submodules={})
78 self._delegate: GeneralizedComponent | None = None
80 def set_original_component(self, component: torch.nn.Module) -> None:
81 super().set_original_component(component)
82 # This subclass's own __init__ takes `name: str` (non-optional), so
83 # self.name is always a str here -- but the base GeneralizedComponent
84 # attribute is typed `str | None`, which is all mypy sees without this
85 # narrowing. MoEBridge/GatedMLPBridge both require a plain `str` name.
86 assert self.name is not None
87 if hasattr(component, "router") and hasattr(component, "experts"):
88 if not isinstance(component.router, torch.nn.Module):
89 raise TypeError(
90 f"{type(component).__name__}.router must be an nn.Module; "
91 f"got {type(component.router).__name__}."
92 )
93 if not isinstance(component.experts, (torch.nn.ModuleList, torch.nn.ModuleDict)):
94 raise TypeError(
95 f"{type(component).__name__}.experts must be a registered "
96 "module collection (nn.ModuleList or nn.ModuleDict); got "
97 f"{type(component.experts).__name__}."
98 )
99 delegate: GeneralizedComponent = MoEBridge(
100 name=self.name,
101 config=self.config,
102 submodules={"gate": LinearBridge(name="router")},
103 )
104 elif hasattr(component, "gate") and hasattr(component, "up") and hasattr(component, "down"):
105 for field in ("gate", "up", "down"):
106 value = getattr(component, field)
107 if not isinstance(value, torch.nn.Module):
108 raise TypeError(
109 f"{type(component).__name__}.{field} must be an "
110 f"nn.Module; got {type(value).__name__}."
111 )
112 delegate = GatedMLPBridge(
113 name=self.name,
114 config=self.config,
115 submodules={
116 "gate": LinearBridge(name="gate"),
117 "in": LinearBridge(name="up"),
118 "out": LinearBridge(name="down"),
119 },
120 )
121 else:
122 raise ValueError(
123 f"Block.mlp is a {type(component).__name__} with neither "
124 "an MoE layer's (router, experts) nor a gated MLP's "
125 "(gate, up, down) attributes -- this adapter doesn't know "
126 "how to wrap it."
127 )
128 delegate.set_original_component(component)
129 # Not `self._delegate = delegate`: normal registration would
130 # duplicate hook_in/hook_out under a nested `._delegate.` path,
131 # risking a broad hook selector firing twice.
132 # `_delegate` is an execution helper, absent from named_modules()
133 # -- safe since parameters/state_dict/dtype are read from the raw
134 # wrapped model (see PretrainModelContainer), not this tree.
135 object.__setattr__(self, "_delegate", delegate)
137 def forward(self, *args: Any, **kwargs: Any) -> Any:
138 assert self._delegate is not None, f"{self.name}: original component not set"
139 if args: 139 ↛ 141line 139 didn't jump to line 141 because the condition on line 139 was always true
140 args = (self.hook_in(args[0]),) + args[1:]
141 elif "hidden_states" in kwargs:
142 kwargs = {**kwargs, "hidden_states": self.hook_in(kwargs["hidden_states"])}
143 output = self._delegate(*args, **kwargs)
145 if isinstance(output, tuple):
146 if len(output) == 0: 146 ↛ 147line 146 didn't jump to line 147 because the condition on line 146 was never true
147 raise TypeError(
148 "DenseOrMoEFeedForwardBridge expected a non-empty tuple "
149 "whose first element is a torch.Tensor"
150 )
152 first = output[0]
154 if not isinstance(first, torch.Tensor): 154 ↛ 155line 154 didn't jump to line 155 because the condition on line 154 was never true
155 raise TypeError(
156 "DenseOrMoEFeedForwardBridge expected the first tuple element "
157 f"to be a torch.Tensor, got {type(first).__name__}"
158 )
160 hooked_first = self.hook_out(first)
162 # Preserve every auxiliary element without sending it through HookPoint.
163 return (hooked_first, *output[1:])
165 if not isinstance(output, torch.Tensor): 165 ↛ 166line 165 didn't jump to line 166 because the condition on line 165 was never true
166 raise TypeError(
167 "DenseOrMoEFeedForwardBridge expected a torch.Tensor or a tuple "
168 f"whose first element is a torch.Tensor, got {type(output).__name__}"
169 )
171 return self.hook_out(output)
174class _LogitsAttrDict(dict):
175 """Makes a plain dict's keys accessible as attributes (`d.logits` reads
176 `d["logits"]`), so a source model's plain-dict forward output satisfies
177 the `hasattr(output, "logits")` contract `TransformerBridge` expects.
178 Behaves as a plain dict everywhere else (indexing, `.get`, `in`, ...).
179 """
181 def __getattr__(self, key: str) -> Any:
182 try:
183 return self[key]
184 except KeyError as e:
185 raise AttributeError(key) from e
188class PretrainModelContainer(torch.nn.Module):
189 """Wraps the source model one level deeper (`container.inner`) so its
190 own `embed`/`blocks` attrs don't collide with `TransformerBridge`
191 component_mapping keys, normalizes the forward return to the `.logits`
192 contract, and strips `_BRIDGE_COMPAT_KWARGS`. Applied automatically by
193 `build_pretrain_bridge`.
194 """
196 def __init__(self, model: torch.nn.Module) -> None:
197 super().__init__()
198 self.inner = model
200 def forward(self, *args: Any, **kwargs: Any) -> Any:
201 filtered = {k: v for k, v in kwargs.items() if k not in _BRIDGE_COMPAT_KWARGS}
202 output = self.inner(*args, **filtered)
204 # Already-normalized or already-HF-style outputs pass through
205 # unchanged, but only after checking .logits is actually a tensor
206 # -- an object merely exposing the attribute isn't enough.
207 if isinstance(output, _LogitsAttrDict):
208 if "logits" not in output:
209 raise ValueError(
210 f"{type(self.inner).__name__}.forward returned a "
211 "_LogitsAttrDict without a 'logits' key."
212 )
213 if not isinstance(output["logits"], torch.Tensor): 213 ↛ 214line 213 didn't jump to line 214 because the condition on line 213 was never true
214 raise TypeError(
215 f"{type(self.inner).__name__}.forward returned a "
216 f"_LogitsAttrDict with a non-tensor 'logits' value: "
217 f"{type(output['logits']).__name__}."
218 )
219 return output
221 # try/except, not hasattr(): hasattr() would evaluate a
222 # property-backed .logits once, then a separate read would
223 # evaluate it again. This reads it exactly once.
224 try:
225 logits = output.logits
226 except AttributeError:
227 pass
228 else:
229 if not isinstance(logits, torch.Tensor):
230 raise TypeError(
231 f"{type(self.inner).__name__}.forward returned a "
232 f"{type(output).__name__} with a non-tensor .logits value: "
233 f"{type(logits).__name__}."
234 )
235 return output
237 if isinstance(output, torch.Tensor):
238 return output
240 # TransformerBridge extracts output[0] as logits for tuple returns.
241 if isinstance(output, tuple):
242 if output and isinstance(output[0], torch.Tensor):
243 return output
244 raise TypeError(
245 f"{type(self.inner).__name__}.forward returned a tuple whose "
246 f"first element is a {type(output[0]).__name__ if output else 'empty tuple'}, "
247 "not a torch.Tensor -- TransformerBridge extracts output[0] as "
248 "logits for tuple returns, so it must be a tensor."
249 )
251 # The primary case: a plain dict, normalized into _LogitsAttrDict.
252 if isinstance(output, dict):
253 if "logits" not in output:
254 raise ValueError(
255 f"{type(self.inner).__name__}.forward returned a dict with keys "
256 # list(), not sorted(): sorted() raises TypeError on
257 # heterogeneous keys, which would mask this error.
258 f"{list(output.keys())}, but PretrainModelContainer requires a "
259 "'logits' key -- without it, TransformerBridge's own "
260 "hasattr(output, 'logits') check would silently fail the same "
261 "way this container exists to prevent."
262 )
263 if not isinstance(output["logits"], torch.Tensor):
264 raise TypeError(
265 f"{type(self.inner).__name__}.forward returned a dict whose "
266 f"'logits' value is a {type(output['logits']).__name__}, not a "
267 "torch.Tensor."
268 )
269 return _LogitsAttrDict(output)
271 raise TypeError(
272 f"{type(self.inner).__name__}.forward must return a torch.Tensor, a "
273 "tuple whose first element is a tensor, a dict containing a "
274 "tensor-valued 'logits' key, or an object with a tensor-valued "
275 f".logits attribute; got {type(output).__name__}."
276 )
279class NativeForwardAttentionBridge(AttentionBridge):
280 """Opaque attention bridge that delegates to the source attention.
282 This adapter intentionally exposes only input/output attention hooks.
283 It has no mapped Q/K/V/O projection components, so the standard
284 per-head aliases and weight aliases do not apply.
285 """
287 hook_aliases = {}
288 property_aliases = {}
289 supports_split_qkv_fork = False
292class PretrainArchitectureAdapter(ArchitectureAdapter):
293 """Adapter for a decoder-only transformer using RoPE, RMSNorm, gated
294 SwiGLU MLPs, and optional sparse MoE feed-forward layers.
296 Uses an opaque `NativeForwardAttentionBridge` (delegates to
297 `Attention.forward`) so RoPE runs under the source's adjacent-pair
298 convention rather than HF's rotate-half -- at the cost of only
299 block-level hooks, no per-head hooks.
300 """
302 def __init__(self, cfg: Any) -> None:
303 super().__init__(cfg)
305 # Also sets uses_rms_norm=True: norm bridges fall back to it when the
306 # wrapped norm's class name doesn't identify itself as RMSNorm, and a
307 # False fallback would mean-center RMS hook intermediates.
308 self._set_rms_rotary_defaults()
310 self.component_mapping = {
311 # "inner." because this adapter expects the source model to
312 # arrive wrapped in `PretrainModelContainer` -- see that
313 # class's docstring for why.
314 "embed": EmbeddingBridge(name="inner.embed"),
315 "blocks": DelegatedAttentionBlockBridge(
316 name="inner.blocks",
317 config=self.cfg,
318 submodules={
319 "ln1": RMSNormalizationBridge(name="norm1", config=self.cfg),
320 "attn": NativeForwardAttentionBridge(
321 name="attn",
322 config=self.cfg,
323 submodules={}, # opaque wrap -- see class docstring
324 ),
325 "ln2": RMSNormalizationBridge(name="norm2", config=self.cfg),
326 "mlp": DenseOrMoEFeedForwardBridge(name="mlp", config=self.cfg),
327 },
328 ),
329 "ln_final": RMSNormalizationBridge(name="inner.norm_f", config=self.cfg),
330 "unembed": UnembeddingBridge(name="inner.lm_head"),
331 }
334def build_pretrain_bridge(
335 model: torch.nn.Module,
336 cfg: TransformerBridgeConfig,
337 *,
338 device: Any = None,
339 dtype: torch.dtype | None = None,
340 model_name: str | None = None,
341) -> TransformerBridge:
342 """Public entry point: wraps `model` in `PretrainModelContainer` and
343 builds a `TransformerBridge` around it. Prefer this over calling
344 `build_bridge_from_module` directly -- the container is easy to forget.
346 `device`/`dtype`/`model_name` forward to `build_bridge_from_module`
347 only when explicitly given.
349 `bridge.train()`/`.eval()` propagate to `model` via
350 `TransformerBridge.train()` itself, which sets mode on
351 `original_model` in addition to the registered module tree
352 (`original_model` is deliberately not a registered submodule, so
353 `nn.Module.train()`'s own recursion never reaches it). This adapter
354 needs nothing extra for mode propagation.
356 Setting mode on `model` directly still works too and stays in sync.
357 """
358 from transformer_lens.model_bridge.sources._bridge_builder import (
359 build_bridge_from_module,
360 )
362 kwargs: dict[str, Any] = {}
363 if device is not None: 363 ↛ 364line 363 didn't jump to line 364 because the condition on line 363 was never true
364 kwargs["device"] = device
365 if dtype is not None: 365 ↛ 366line 365 didn't jump to line 366 because the condition on line 365 was never true
366 kwargs["dtype"] = dtype
367 if model_name is not None:
368 kwargs["model_name"] = model_name
370 return build_bridge_from_module(
371 PretrainModelContainer(model),
372 architecture=ARCHITECTURE_NAME,
373 tl_config=cfg,
374 **kwargs,
375 )