Coverage for transformer_lens/model_bridge/supported_architectures/pretrain.py: 89%
112 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-08-11 18:50 +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.bridge import TransformerBridge
46from transformer_lens.model_bridge.generalized_components import (
47 AttentionBridge,
48 DelegatedAttentionBlockBridge,
49 EmbeddingBridge,
50 GatedMLPBridge,
51 LinearBridge,
52 MoEBridge,
53 RMSNormalizationBridge,
54 UnembeddingBridge,
55)
56from transformer_lens.model_bridge.generalized_components.base import (
57 GeneralizedComponent,
58)
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 either a dense SwiGLU MLP or a sparse MoE layer behind a
71 common interface. Dispatch is determined by structural inspection
72 (`router`/`experts` vs `gate`/`up`/`down`) rather than configuration,
73 so dense, MoE, and mixed architectures all use the same component
74 mapping. This identifies the supported protocol -- it does not
75 validate that a module merely sharing those attribute names actually
76 implements the matching forward behavior.
78 `hasattr` alone would let a module with, say, both `router`/`experts`
79 and `gate`/`up`/`down` (or `router`/`experts` of the wrong types) win
80 the MoE branch by attribute-name coincidence and fail later with a
81 confusing error from deep inside `MoEBridge`, or not fail until
82 forward time. `set_original_component` therefore also checks the
83 basic shape of whichever protocol wins: `router`/`gate`/`up`/`down`
84 must themselves be modules, and `experts` must be a *registered*
85 module collection (`nn.ModuleList`/`nn.ModuleDict`) -- a plain Python
86 list/tuple of `nn.Module` experts is rejected even though every
87 element is itself a valid module, because modules held in an
88 ordinary list aren't registered as children and would silently drop
89 out of `parameters()`/`state_dict()`/`.to(...)`/`train()`/`eval()`,
90 contradicting this adapter's lifecycle guarantees.
91 """
93 def __init__(self, name: str, config: Any):
94 super().__init__(name, config=config, submodules={})
95 self._delegate: GeneralizedComponent | None = None
97 def set_original_component(self, component: torch.nn.Module) -> None:
98 super().set_original_component(component)
99 # This subclass's own __init__ takes `name: str` (non-optional), so
100 # self.name is always a str here -- but the base GeneralizedComponent
101 # attribute is typed `str | None`, which is all mypy sees without this
102 # narrowing. MoEBridge/GatedMLPBridge both require a plain `str` name.
103 assert self.name is not None
104 if hasattr(component, "router") and hasattr(component, "experts"):
105 if not isinstance(component.router, torch.nn.Module):
106 raise TypeError(
107 f"{type(component).__name__}.router must be an nn.Module; "
108 f"got {type(component.router).__name__}."
109 )
110 if not isinstance(component.experts, (torch.nn.ModuleList, torch.nn.ModuleDict)):
111 raise TypeError(
112 f"{type(component).__name__}.experts must be a registered "
113 "module collection (nn.ModuleList or nn.ModuleDict); got "
114 f"{type(component.experts).__name__}."
115 )
116 delegate: GeneralizedComponent = MoEBridge(
117 name=self.name,
118 config=self.config,
119 submodules={"gate": LinearBridge(name="router")},
120 )
121 elif hasattr(component, "gate") and hasattr(component, "up") and hasattr(component, "down"):
122 for field in ("gate", "up", "down"):
123 value = getattr(component, field)
124 if not isinstance(value, torch.nn.Module):
125 raise TypeError(
126 f"{type(component).__name__}.{field} must be an "
127 f"nn.Module; got {type(value).__name__}."
128 )
129 delegate = GatedMLPBridge(
130 name=self.name,
131 config=self.config,
132 submodules={
133 "gate": LinearBridge(name="gate"),
134 "in": LinearBridge(name="up"),
135 "out": LinearBridge(name="down"),
136 },
137 )
138 else:
139 raise ValueError(
140 f"Block.mlp is a {type(component).__name__} with neither "
141 "an MoE layer's (router, experts) nor a gated MLP's "
142 "(gate, up, down) attributes -- this adapter doesn't know "
143 "how to wrap it."
144 )
145 delegate.set_original_component(component)
146 # Not `self._delegate = delegate`: normal registration would
147 # duplicate hook_in/hook_out under a nested `._delegate.` path,
148 # risking a broad hook selector firing twice.
149 # `_delegate` is an execution helper, absent from named_modules()
150 # -- safe since parameters/state_dict/dtype are read from the raw
151 # wrapped model (see PretrainModelContainer), not this tree.
152 object.__setattr__(self, "_delegate", delegate)
154 def forward(self, *args: Any, **kwargs: Any) -> Any:
155 assert self._delegate is not None, f"{self.name}: original component not set"
156 if args: 156 ↛ 158line 156 didn't jump to line 158 because the condition on line 156 was always true
157 args = (self.hook_in(args[0]),) + args[1:]
158 elif "hidden_states" in kwargs:
159 kwargs = {**kwargs, "hidden_states": self.hook_in(kwargs["hidden_states"])}
160 output = self._delegate(*args, **kwargs)
162 if isinstance(output, tuple):
163 if len(output) == 0: 163 ↛ 164line 163 didn't jump to line 164 because the condition on line 163 was never true
164 raise TypeError(
165 "DenseOrMoEFeedForwardBridge expected a non-empty tuple "
166 "whose first element is a torch.Tensor"
167 )
169 first = output[0]
171 if not isinstance(first, torch.Tensor): 171 ↛ 172line 171 didn't jump to line 172 because the condition on line 171 was never true
172 raise TypeError(
173 "DenseOrMoEFeedForwardBridge expected the first tuple element "
174 f"to be a torch.Tensor, got {type(first).__name__}"
175 )
177 hooked_first = self.hook_out(first)
179 # Preserve every auxiliary element without sending it through HookPoint.
180 return (hooked_first, *output[1:])
182 if not isinstance(output, torch.Tensor): 182 ↛ 183line 182 didn't jump to line 183 because the condition on line 182 was never true
183 raise TypeError(
184 "DenseOrMoEFeedForwardBridge expected a torch.Tensor or a tuple "
185 f"whose first element is a torch.Tensor, got {type(output).__name__}"
186 )
188 return self.hook_out(output)
191class _LogitsAttrDict(dict):
192 """Makes a plain dict's keys accessible as attributes (`d.logits` reads
193 `d["logits"]`), so a source model's plain-dict forward output satisfies
194 the `hasattr(output, "logits")` contract `TransformerBridge` expects.
195 Behaves as a plain dict everywhere else (indexing, `.get`, `in`, ...).
196 """
198 def __getattr__(self, key: str) -> Any:
199 try:
200 return self[key]
201 except KeyError as e:
202 raise AttributeError(key) from e
205class PretrainModelContainer(torch.nn.Module):
206 """Internal detail -- `build_pretrain_bridge` applies this
207 automatically. Three responsibilities, all in this container's own
208 `forward`:
210 1. Avoids a `TransformerBridge.__getattr__` collision: a source
211 model's own `self.embed`/`self.blocks` clashes with identically
212 named component_mapping keys. Wrapping one level deeper
213 (`container.inner.embed`) fixes this without touching the source
214 model.
215 2. Normalizes the return value to the `.logits` contract
216 `TransformerBridge` expects: a plain `"logits"` dict (the target
217 architecture's actual shape) is wrapped in `_LogitsAttrDict`; a bare
218 tensor, tensor-first tuple, or object already exposing `.logits`
219 passes through after validating the tensor is present and is a
220 tensor; anything else raises immediately with a clear message.
221 3. Strips `_BRIDGE_COMPAT_KWARGS` from kwargs before calling the
222 wrapped model (see that constant's comment).
224 Also sidesteps a `BlockBridge` convention where a bare-tensor block
225 output gets wrapped in a 1-tuple for "standalone hidden_states calls":
226 since this target's blocks take `cos`/`sin` too, that path never
227 triggers, so the source forward loop needs no changes to be bridged.
229 `self.inner` is a regular registered submodule, so
230 `container.train()`/`.eval()` already recurse into it via the normal
231 `nn.Module` traversal -- no override needed here. The propagation gap
232 lives one level up, at `TransformerBridge` itself (see
233 `build_pretrain_bridge`), whose `.train()`/`.eval()` do not walk down
234 to `original_model`.
235 """
237 def __init__(self, model: torch.nn.Module) -> None:
238 super().__init__()
239 self.inner = model
241 def forward(self, *args: Any, **kwargs: Any) -> Any:
242 filtered = {k: v for k, v in kwargs.items() if k not in _BRIDGE_COMPAT_KWARGS}
243 output = self.inner(*args, **filtered)
245 # Already-normalized or already-HF-style outputs pass through
246 # unchanged, but only after checking .logits is actually a tensor
247 # -- an object merely exposing the attribute isn't enough.
248 if isinstance(output, _LogitsAttrDict):
249 if "logits" not in output:
250 raise ValueError(
251 f"{type(self.inner).__name__}.forward returned a "
252 "_LogitsAttrDict without a 'logits' key."
253 )
254 if not isinstance(output["logits"], torch.Tensor): 254 ↛ 255line 254 didn't jump to line 255 because the condition on line 254 was never true
255 raise TypeError(
256 f"{type(self.inner).__name__}.forward returned a "
257 f"_LogitsAttrDict with a non-tensor 'logits' value: "
258 f"{type(output['logits']).__name__}."
259 )
260 return output
262 # try/except, not hasattr(): hasattr() would evaluate a
263 # property-backed .logits once, then a separate read would
264 # evaluate it again. This reads it exactly once.
265 try:
266 logits = output.logits
267 except AttributeError:
268 pass
269 else:
270 if not isinstance(logits, torch.Tensor):
271 raise TypeError(
272 f"{type(self.inner).__name__}.forward returned a "
273 f"{type(output).__name__} with a non-tensor .logits value: "
274 f"{type(logits).__name__}."
275 )
276 return output
278 if isinstance(output, torch.Tensor):
279 return output
281 # TransformerBridge extracts output[0] as logits for tuple returns.
282 if isinstance(output, tuple):
283 if output and isinstance(output[0], torch.Tensor):
284 return output
285 raise TypeError(
286 f"{type(self.inner).__name__}.forward returned a tuple whose "
287 f"first element is a {type(output[0]).__name__ if output else 'empty tuple'}, "
288 "not a torch.Tensor -- TransformerBridge extracts output[0] as "
289 "logits for tuple returns, so it must be a tensor."
290 )
292 # The primary case: a plain dict, normalized into _LogitsAttrDict.
293 if isinstance(output, dict):
294 if "logits" not in output:
295 raise ValueError(
296 f"{type(self.inner).__name__}.forward returned a dict with keys "
297 # list(), not sorted(): sorted() raises TypeError on
298 # heterogeneous keys, which would mask this error.
299 f"{list(output.keys())}, but PretrainModelContainer requires a "
300 "'logits' key -- without it, TransformerBridge's own "
301 "hasattr(output, 'logits') check would silently fail the same "
302 "way this container exists to prevent."
303 )
304 if not isinstance(output["logits"], torch.Tensor):
305 raise TypeError(
306 f"{type(self.inner).__name__}.forward returned a dict whose "
307 f"'logits' value is a {type(output['logits']).__name__}, not a "
308 "torch.Tensor."
309 )
310 return _LogitsAttrDict(output)
312 raise TypeError(
313 f"{type(self.inner).__name__}.forward must return a torch.Tensor, a "
314 "tuple whose first element is a tensor, a dict containing a "
315 "tensor-valued 'logits' key, or an object with a tensor-valued "
316 f".logits attribute; got {type(output).__name__}."
317 )
320class NativeForwardAttentionBridge(AttentionBridge):
321 """Opaque attention bridge that delegates to the source attention.
323 This adapter intentionally exposes only input/output attention hooks.
324 It has no mapped Q/K/V/O projection components, so the standard
325 per-head aliases and weight aliases do not apply.
326 """
328 hook_aliases = {}
329 property_aliases = {}
330 supports_split_qkv_fork = False
333class PretrainArchitectureAdapter(ArchitectureAdapter):
334 """Adapter for a decoder-only transformer using RoPE, RMSNorm, gated
335 SwiGLU MLPs, and optional sparse MoE feed-forward layers.
337 Uses an opaque `NativeForwardAttentionBridge` with no attention
338 projection submodules, not `JointQKVAttentionBridge`/
339 `PositionEmbeddingsAttentionBridge`: those reimplement RoPE via HF's
340 rotate-half convention, wrong for a source model using the
341 adjacent-pair convention. The opaque bridge delegates unchanged to
342 `Attention.forward`, so RoPE runs as written -- at the cost of no
343 per-head hooks, only block-level
344 `resid_pre`/`resid_mid`/`resid_post`.
346 Blocks use `DelegatedAttentionBlockBridge` rather than plain
347 `BlockBridge`: that existing abstraction already exists for
348 architectures where attention is delegated wholesale and the
349 split-qkv-fork block-level aliases (`hook_attn_in`/`hook_q_input`/
350 `hook_k_input`/`hook_v_input`) don't apply. It complements
351 `NativeForwardAttentionBridge.supports_split_qkv_fork = False` (which
352 prevents the split-QKV-fork machinery and its associated HookPoints
353 from being exposed for this attention component) by also removing the
354 now-dangling block-level aliases that would otherwise point at them.
355 `hook_attn_out` is untouched by either change, since the attention
356 component still fires its own `hook_out` normally.
358 `self.cfg` is mutated in place, not copied (matches `nanogpt.py`'s
359 convention) -- callers holding another reference to the same config
360 will see these fields change.
362 Bridges built through `build_pretrain_bridge` are given a
363 mode-propagating subclass so `.train()`/`.eval()` reach the wrapped
364 source model (see that function's docstring) -- this adapter class
365 itself has no lifecycle behavior of its own.
366 """
368 def __init__(self, cfg: Any) -> None:
369 super().__init__(cfg)
371 self.cfg.normalization_type = "RMS"
372 self.cfg.positional_embedding_type = "rotary"
373 self.cfg.final_rms = True
374 self.cfg.gated_mlp = True
375 self.cfg.attn_only = False
377 self.component_mapping = {
378 # "inner." because this adapter expects the source model to
379 # arrive wrapped in `PretrainModelContainer` -- see that
380 # class's docstring for why.
381 "embed": EmbeddingBridge(name="inner.embed"),
382 "blocks": DelegatedAttentionBlockBridge(
383 name="inner.blocks",
384 config=self.cfg,
385 submodules={
386 "ln1": RMSNormalizationBridge(name="norm1", config=self.cfg),
387 "attn": NativeForwardAttentionBridge(
388 name="attn",
389 config=self.cfg,
390 submodules={}, # opaque wrap -- see class docstring
391 ),
392 "ln2": RMSNormalizationBridge(name="norm2", config=self.cfg),
393 "mlp": DenseOrMoEFeedForwardBridge(name="mlp", config=self.cfg),
394 },
395 ),
396 "ln_final": RMSNormalizationBridge(name="inner.norm_f", config=self.cfg),
397 "unembed": UnembeddingBridge(name="inner.lm_head"),
398 }
401def build_pretrain_bridge(
402 model: torch.nn.Module,
403 cfg: TransformerBridgeConfig,
404 *,
405 device: Any = None,
406 dtype: torch.dtype | None = None,
407 model_name: str | None = None,
408) -> TransformerBridge:
409 """Public entry point: wraps `model` in `PretrainModelContainer` and
410 builds a `TransformerBridge` around it. Prefer this over calling
411 `build_bridge_from_module` directly -- the container is easy to forget.
413 `device`/`dtype`/`model_name` forward to `build_bridge_from_module`
414 only when explicitly given.
416 `bridge.train()`/`.eval()` propagate to `model` via
417 `TransformerBridge.train()` itself, which sets mode on
418 `original_model` in addition to the registered module tree
419 (`original_model` is deliberately not a registered submodule, so
420 `nn.Module.train()`'s own recursion never reaches it). This adapter
421 needs nothing extra for mode propagation.
423 Setting mode on `model` directly still works too and stays in sync.
424 """
425 from transformer_lens.model_bridge.sources._bridge_builder import (
426 build_bridge_from_module,
427 )
429 kwargs: dict[str, Any] = {}
430 if device is not None: 430 ↛ 431line 430 didn't jump to line 431 because the condition on line 430 was never true
431 kwargs["device"] = device
432 if dtype is not None: 432 ↛ 433line 432 didn't jump to line 433 because the condition on line 432 was never true
433 kwargs["dtype"] = dtype
434 if model_name is not None:
435 kwargs["model_name"] = model_name
437 return build_bridge_from_module(
438 PretrainModelContainer(model),
439 architecture=ARCHITECTURE_NAME,
440 tl_config=cfg,
441 **kwargs,
442 )