Coverage for transformer_lens/model_bridge/supported_architectures/raven.py: 49%
53 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-08-11 18:50 +0000
1"""Raven / Huginn architecture adapter (RavenForCausalLM).
3Model family: tomg-group-umd/huginn-0125 ("Huginn"), a depth-recurrent
4("latent reasoning") decoder from Geiping et al. Loaded via remote code
5(``auto_map`` → ``raven_modeling_minimal.RavenForCausalLM``), so
6``trust_remote_code=True`` is required.
8Architecture overview
9---------------------
10Huginn is NOT a flat stack of transformer layers. Its forward has three
11phases, all operating on the same residual width ``n_embd`` (5280):
13 wte → prelude (P physical blocks)
14 → [ recurrent core: R physical blocks, applied N times ]
15 → coda (C physical blocks) → ln_f → lm_head
17with P = ``n_layers_in_prelude`` (2), R = ``n_layers_in_recurrent_block``
18(4), C = ``n_layers_in_coda`` (2). ``num_hidden_layers`` (8) counts the
19*physical* blocks (2 + 4 + 2), stored as three separate ``ModuleList``s under
20``model.transformer`` (``prelude`` / ``core_block`` / ``coda``), NOT one
21``model.layers``.
23The recurrence is the defining feature. ``RavenForCausalLM.forward`` calls
24``iterate_forward`` → ``core_block_forward``, which runs the SAME four
25``core_block`` modules N times. Each step re-injects the prelude output:
27 x = adapter(cat([latent_state, prelude_output], dim=-1)) # injection
28 for block in core_block: x = block(x) # 4 blocks
30``N`` (``num_steps``) is a RUNTIME argument to ``forward``, not a fixed
31config value. At eval it defaults to ``config.mean_recurrence`` (32) via
32``randomized_iteration_sampler``; a caller may pass any ``num_steps``.
34Each block is a ``SandwichBlock``: four RMSNorms with post-residual
35normalisation (``x = norm_2(attn(norm_1(x)) + x)``; ``x = norm_4(mlp(norm_3
36(x)) + x)``) — the residual stream itself is renormalised after each add,
37unlike a standard pre-norm transformer. Attention is MHA (55 heads == 55 kv
38heads) with a COMBINED ``Wqkv`` projection plus a learned additive ``qk_bias``
39parameter and RoPE (base 50000). The MLP is a gated SiLU MLP with a combined
40gate+up ``fc`` projection. Embeddings are scaled by ``√n_embd`` (≈72.66) in
41the HF forward, and ``lm_head`` is tied to ``wte``.
43Key adapter decisions
44---------------------
451. Full delegation, like Ouro. The recurrence, the prelude re-injection, the
46 emb-scale, the sandwich norms and RoPE all live inside the remote-code
47 ``forward`` that the bridge delegates to, so a single forward pass is
48 numerically correct with no loop handling here.
502. ``OpaqueBlockBridge`` for all three block lists (``prelude`` / ``core_block``
51 / ``coda``): ``BlockBridge``'s hook aliases hardcode a standard pre-norm flow
52 (hook_resid_mid, ln1→attn→ln2→mlp) that the post-residual ``SandwichBlock``
53 does not follow. ``OpaqueBlockBridge`` delegates the whole block and exposes
54 only ``hook_in`` / ``hook_out`` on the residual stream, which is correct
55 regardless of the internal norm placement. Each block's inner attn / mlp /
56 norms are still declared as submodules so they wrap the live HF modules and
57 their hooks fire.
593. Recurrent core hooks. Because the core loop lives inside the HF forward,
60 the four ``core_block`` blocks' ``hook_in`` / ``hook_out`` fire once PER
61 recurrence step (N times per forward), and ``run_with_cache`` keeps the
62 final step. This is the load-bearing behavioural difference from a flat
63 decoder and mirrors Ouro's looped-depth semantics. Separately addressing
64 an individual recurrence step (e.g. logit-lens across steps) is NOT
65 expressible through the static ``core_block.{i}.hook_out`` names — it
66 requires the model's native ``iterate_one_step`` / ``predict_from_latents``
67 interface. Deliberately not mapped: ``transformer.adapter`` (the injection
68 Linear) and ``HuginnDynamicCache``'s block-indexed slot layout.
704. ``applicable_phases = []``. Huginn diverges too far from the transformer-
71 shaped ``verify_models`` phases to score meaningfully: post-residual
72 sandwich norms, a runtime recurrence count, combined QKV + ``qk_bias``,
73 and — decisively — a RANDOM initial latent state (``initialize_state``
74 uses ``torch.randn_like``), which makes the forward non-deterministic
75 across calls unless the RNG is seeded or ``input_states`` is supplied.
76 Correctness is covered in the integration tests, where the seed is pinned
77 before both the bridge and HF calls (same decision spirit as nemotron_h /
78 mamba). ``ln_f`` is applied mid-network (after the recurrence, feeding the
79 coda) as well as at the end, so ``supports_fold_ln = False`` — folding it
80 into ``W_U`` would corrupt the coda input.
82Config propagation
83------------------
84The recurrence-shape attributes (``mean_recurrence``, ``n_layers_in_prelude``
85/ ``_recurrent_block`` / ``_coda``, ``mean_backprop_depth``, ``injection_type``,
86``qk_bias``) are surfaced on ``self.cfg`` here AND added to the two
87``_HF_PASSTHROUGH_ATTRS`` lists so analysis tooling can read them off a booted
88bridge.
90Remote-code loading (transformers v5)
91------------------------------------
92Huginn's remote code targets transformers 4.44 and breaks under v5 (5.8.1) in
93two ways that ``prepare_loading`` patches: (1) ``_tied_weights_keys`` is a list,
94but v5's ``tie_weights`` expects a dict, so the model does not even construct;
95(2) v5's meta-device load re-invokes ``PreTrainedModel._init_weights`` on
96already-materialised modules and would re-randomise the checkpoint. See
97``prepare_loading`` for both. (Huginn does not use ``ROPE_INIT_FUNCTIONS``; it
98precomputes its own ``freqs_cis``, so no RoPE patch is needed.)
100Optional parameters (may be absent from a state_dict)
101----------------------------------------------------
102Huginn has ``bias=False`` everywhere except the attention ``qk_bias``
103parameter; RMSNorm has no bias. Weight processing must tolerate missing
104biases via ``ProcessWeights._safe_get_tensor()``.
105"""
107import sys
108from typing import Any
110from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter
111from transformer_lens.model_bridge.generalized_components import (
112 EmbeddingBridge,
113 LinearBridge,
114 MLPBridge,
115 OpaqueBlockBridge,
116 RMSNormalizationBridge,
117 UnembeddingBridge,
118)
119from transformer_lens.model_bridge.generalized_components.attention import (
120 AttentionBridge,
121)
124class RavenArchitectureAdapter(ArchitectureAdapter):
125 """Architecture adapter for RavenForCausalLM (Huginn depth-recurrent decoder).
127 Prelude / weight-tied recurrent core / coda phases over a shared residual
128 width. The recurrence and prelude re-injection live inside the remote-code
129 HF forward, which the bridge delegates to; see the module docstring for the
130 full set of adapter decisions.
131 """
133 # Huginn is off the transformer-shaped verify_models path: post-residual
134 # sandwich norms, runtime recurrence count, and a random initial latent
135 # state make the phases non-meaningful. Correctness lives in the
136 # integration tests (seed pinned before bridge and HF calls).
137 applicable_phases: list[int] = []
139 def __init__(self, cfg: Any) -> None:
140 """Initialize the Raven / Huginn architecture adapter."""
141 super().__init__(cfg)
143 # Standard weight-processing / norm flags.
144 self.cfg.normalization_type = "RMS"
145 self.cfg.uses_rms_norm = True
146 self.cfg.positional_embedding_type = "rotary"
147 self.cfg.final_rms = True
148 self.cfg.gated_mlp = True
149 self.cfg.attn_only = False
151 # ln_f (transformer.ln_f) is applied after the recurrence (feeding the
152 # coda) AND at the very end, so it is not a final-only norm. Folding it
153 # into W_U would corrupt the coda's input.
154 self.supports_fold_ln = False
156 # Surface the recurrence-shape attributes on cfg so they are present on
157 # both the HF-boot path (also via _HF_PASSTHROUGH_ATTRS) and the
158 # synthetic-config path used by the unit tests.
159 setattr(self.cfg, "mean_recurrence", getattr(cfg, "mean_recurrence", 32))
160 setattr(self.cfg, "mean_backprop_depth", getattr(cfg, "mean_backprop_depth", 8))
161 setattr(self.cfg, "n_layers_in_prelude", getattr(cfg, "n_layers_in_prelude", 2))
162 setattr(
163 self.cfg, "n_layers_in_recurrent_block", getattr(cfg, "n_layers_in_recurrent_block", 4)
164 )
165 setattr(self.cfg, "n_layers_in_coda", getattr(cfg, "n_layers_in_coda", 2))
166 setattr(self.cfg, "injection_type", getattr(cfg, "injection_type", "linear"))
167 setattr(self.cfg, "qk_bias", getattr(cfg, "qk_bias", True))
169 # Full delegation to the HF forward — no HT-format weight reshaping.
170 self.weight_processing_conversions = {}
172 self.component_mapping = {
173 "embed": EmbeddingBridge(name="transformer.wte"),
174 # Three separate physical block lists. Each uses OpaqueBlockBridge so
175 # the delegated SandwichBlock forward keeps its post-residual norm
176 # placement while hook_in / hook_out wrap the residual stream. Fresh
177 # submodule instances per list (they bind to distinct HF modules).
178 "prelude": OpaqueBlockBridge(
179 name="transformer.prelude",
180 submodules=self._sandwich_submodules(),
181 ),
182 "core_block": OpaqueBlockBridge(
183 name="transformer.core_block",
184 submodules=self._sandwich_submodules(),
185 ),
186 "coda": OpaqueBlockBridge(
187 name="transformer.coda",
188 submodules=self._sandwich_submodules(),
189 ),
190 "ln_final": RMSNormalizationBridge(name="transformer.ln_f", config=self.cfg),
191 "unembed": UnembeddingBridge(name="lm_head"),
192 }
194 def _sandwich_submodules(self) -> dict[str, Any]:
195 """Build a fresh set of SandwichBlock submodule bridges.
197 Returns new instances on every call so each of the three block lists
198 wraps its own live HF modules rather than sharing bridge objects.
200 Submodule keys mirror the HF attribute names (``norm_1``..``norm_4``,
201 ``attn``, ``mlp``) so weight-key translation is identity. Attention is
202 native (combined ``Wqkv`` + ``qk_bias``, RoPE, custom SDPA path), so it
203 is delegated via ``maintain_native_attention``; only the combined
204 ``qkv`` and output ``o`` projections are exposed. The gated MLP is
205 likewise delegated with its combined gate+up ``fc`` and output ``proj``.
206 """
207 attn = AttentionBridge(
208 name="attn",
209 config=self.cfg,
210 submodules={
211 "qkv": LinearBridge(name="Wqkv"),
212 "o": LinearBridge(name="proj"),
213 },
214 maintain_native_attention=True,
215 requires_attention_mask=True,
216 )
217 # Raven uses a combined Wqkv projection; no separate q/k/v submodules exist.
218 # Strip the default hook_q/hook_k/hook_v aliases so they don't appear as
219 # dead (unresolvable) aliases in the hook-alias resolution audit.
220 attn.hook_aliases = {
221 k: v
222 for k, v in AttentionBridge.hook_aliases.items()
223 if k not in {"hook_q", "hook_k", "hook_v"}
224 }
225 return {
226 "norm_1": RMSNormalizationBridge(name="norm_1", config=self.cfg),
227 "attn": attn,
228 "norm_2": RMSNormalizationBridge(name="norm_2", config=self.cfg),
229 "norm_3": RMSNormalizationBridge(name="norm_3", config=self.cfg),
230 "mlp": MLPBridge(
231 name="mlp",
232 config=self.cfg,
233 submodules={
234 "in": LinearBridge(name="fc"),
235 "out": LinearBridge(name="proj"),
236 },
237 ),
238 "norm_4": RMSNormalizationBridge(name="norm_4", config=self.cfg),
239 }
241 def prepare_loading(self, model_name: str, model_kwargs: dict) -> None:
242 """Patch Huginn's remote code for transformers v5 compatibility.
244 Huginn's modeling code targets transformers 4.44; two things break under
245 v5 (5.8.1), so two patches:
247 1. Tied-weights format. ``RavenForCausalLM._tied_weights_keys`` is a list
248 (``["lm_head.weight"]``, the 4.x format), but v5's ``tie_weights`` ->
249 ``get_expanded_tied_weights_keys`` calls ``.keys()`` on it and raises
250 ``AttributeError``. The model does not even construct. Rewrite it to
251 the v5 dict form ``{"lm_head.weight": "transformer.wte.weight"}``
252 (Huginn ties ``lm_head`` to ``transformer.wte``).
254 2. Weight re-init. Under v5's meta-device load-then-materialise flow,
255 ``PreTrainedModel._init_weights`` is invoked on modules that already
256 hold checkpoint weights, re-randomising them. Guard it to skip modules
257 whose parameters are already on a real (non-meta) device — the same
258 defensive patch openelm.py applies.
260 Args:
261 model_name: The HuggingFace model name/path.
262 model_kwargs: The kwargs dict for from_pretrained().
263 """
264 # Force-import the modeling module so it appears in sys.modules to patch.
265 try:
266 from transformers.dynamic_module_utils import get_class_from_dynamic_module
268 get_class_from_dynamic_module(
269 "raven_modeling_minimal.RavenForCausalLM",
270 model_name,
271 )
272 except Exception:
273 return
275 # Each checkpoint revision gets its own module in sys.modules; patch all.
276 for key in list(sys.modules.keys()):
277 if "raven" not in key.lower() or "modeling" not in key.lower():
278 continue
279 module = sys.modules[key]
281 # Patch 1: tied-weights keys list -> v5 dict form.
282 causal_lm_class = getattr(module, "RavenForCausalLM", None)
283 if causal_lm_class is not None and isinstance(
284 getattr(causal_lm_class, "_tied_weights_keys", None), list
285 ):
286 causal_lm_class._tied_weights_keys = {"lm_head.weight": "transformer.wte.weight"}
288 # Patch 2: don't re-randomise already-loaded weights.
289 pretrained_class = getattr(module, "RavenPreTrainedModel", None)
290 if pretrained_class is None or getattr(pretrained_class, "_tl_patched", False):
291 continue
292 original_init_weights = pretrained_class._init_weights
294 def safe_init_weights(self, mod, _original=original_init_weights):
295 # Only initialise modules still on meta device (pre-loading);
296 # never re-randomise weights already read from the checkpoint.
297 first_param = next(mod.parameters(), None)
298 if first_param is not None and first_param.device.type != "meta":
299 return
300 _original(self, mod)
302 pretrained_class._init_weights = safe_init_weights
303 pretrained_class._tl_patched = True