Coverage for transformer_lens/model_bridge/supported_architectures/raven.py: 92%

1"""Raven / Huginn architecture adapter (RavenForCausalLM). 

2 

3Family ``tomg-group-umd/huginn-0125``: depth-recurrent ("latent reasoning") 

4decoder loaded via remote code. Three phases over one residual width — 

5prelude blocks, a weight-tied recurrent core applied N times (``num_steps`` 

6is a RUNTIME forward argument, defaulting to ``config.mean_recurrence``), 

7then coda blocks. Core blocks are post-residual ``SandwichBlock`` modules 

8(residual renormalised after each add), MHA with combined ``Wqkv`` plus a 

9learned additive ``qk_bias``. 

10 

11Adapter decisions: 

12- Full delegation: recurrence, prelude re-injection, emb-scale, sandwich 

13 norms, and RoPE all run inside the remote-code forward. 

14- ``OpaqueBlockBridge`` for all three block lists: BlockBridge's hook aliases 

15 hardcode the standard pre-norm flow the SandwichBlock does not follow. 

16- Core-block ``hook_in``/``hook_out`` fire once PER recurrence step (N times 

17 per forward); ``run_with_cache`` keeps the final step. Per-step access is 

18 not expressible through the static hook names — it needs the model's native 

19 ``iterate_one_step``/``predict_from_latents`` interface. Deliberately not 

20 mapped: ``transformer.adapter`` and ``HuginnDynamicCache``'s slot layout. 

21- ``applicable_phases = []``: ``initialize_state`` uses ``torch.randn_like``, 

22 so the forward is non-deterministic unless seeded — integration tests pin 

23 the seed around both the bridge and HF calls. 

24- Only bias anywhere is ``qk_bias``; weight processing must tolerate missing 

25 biases via ``ProcessWeights._safe_get_tensor()``. 

26""" 

27 

28from typing import Any 

29 

30from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter 

31from transformer_lens.model_bridge.generalized_components import ( 

32 EmbeddingBridge, 

33 LinearBridge, 

34 OpaqueBlockBridge, 

35 RMSNormalizationBridge, 

36 UnembeddingBridge, 

37) 

38from transformer_lens.model_bridge.generalized_components.attention import ( 

39 AttentionBridge, 

40) 

41from transformer_lens.model_bridge.supported_architectures._remote_code_compat import ( 

42 force_import_remote_class, 

43 iter_remote_modeling_modules, 

44 patch_init_weights_skip_loaded, 

45 retie_weights_keys_v5, 

46) 

47 

48 

49class RavenArchitectureAdapter(ArchitectureAdapter): 

50 """Architecture adapter for RavenForCausalLM (Huginn depth-recurrent decoder). 

51 

52 Prelude / weight-tied recurrent core / coda phases over a shared residual 

53 width. The recurrence and prelude re-injection live inside the remote-code 

54 HF forward, which the bridge delegates to; see the module docstring for the 

55 full set of adapter decisions. 

56 """ 

57 

58 # Huginn is off the transformer-shaped verify_models path: post-residual 

59 # sandwich norms, runtime recurrence count, and a random initial latent 

60 # state make the phases non-meaningful. Correctness lives in the 

61 # integration tests (seed pinned before bridge and HF calls). 

62 applicable_phases: list[int] = [] 

63 # HuginnDynamicCache decode freezes K/V from earlier stochastic calls (fresh 

64 # randn latent per forward), so cached logits can't seed-match the full HF 

65 # forward we verify against; full-prefix recompute per step is exact. 

66 supports_kv_cache: bool = False 

67 

68 # The remote forward ignores attention_mask (its compile_mask call is 

69 # commented out), so left-padded rows silently attend to pad tokens. 

70 supports_batched_generation: bool = False 

71 

72 def __init__(self, cfg: Any) -> None: 

73 """Initialize the Raven / Huginn architecture adapter.""" 

74 super().__init__(cfg) 

75 

76 self._set_rms_rotary_defaults() 

77 

78 # ln_f (transformer.ln_f) is applied after the recurrence (feeding the 

79 # coda) AND at the very end, so it is not a final-only norm. Folding it 

80 # into W_U would corrupt the coda's input. 

81 self.supports_fold_ln = False 

82 

83 # Surface the recurrence-shape attributes on cfg so they are present on 

84 # both the HF-boot path (also via _HF_PASSTHROUGH_ATTRS) and the 

85 # synthetic-config path used by the unit tests. 

86 setattr(self.cfg, "mean_recurrence", getattr(cfg, "mean_recurrence", 32)) 

87 setattr(self.cfg, "mean_backprop_depth", getattr(cfg, "mean_backprop_depth", 8)) 

88 setattr(self.cfg, "n_layers_in_prelude", getattr(cfg, "n_layers_in_prelude", 2)) 

89 setattr( 

90 self.cfg, "n_layers_in_recurrent_block", getattr(cfg, "n_layers_in_recurrent_block", 4) 

91 ) 

92 setattr(self.cfg, "n_layers_in_coda", getattr(cfg, "n_layers_in_coda", 2)) 

93 setattr(self.cfg, "injection_type", getattr(cfg, "injection_type", "linear")) 

94 setattr(self.cfg, "qk_bias", getattr(cfg, "qk_bias", True)) 

95 

96 # Full delegation to the HF forward — no HT-format weight reshaping. 

97 self.weight_processing_conversions = {} 

98 

99 self.component_mapping = { 

100 "embed": EmbeddingBridge(name="transformer.wte"), 

101 # Three separate physical block lists. Each uses OpaqueBlockBridge so 

102 # the delegated SandwichBlock forward keeps its post-residual norm 

103 # placement while hook_in / hook_out wrap the residual stream. Fresh 

104 # submodule instances per list (they bind to distinct HF modules). 

105 "prelude": OpaqueBlockBridge( 

106 name="transformer.prelude", 

107 submodules=self._sandwich_submodules(), 

108 ), 

109 "core_block": OpaqueBlockBridge( 

110 name="transformer.core_block", 

111 submodules=self._sandwich_submodules(), 

112 ), 

113 "coda": OpaqueBlockBridge( 

114 name="transformer.coda", 

115 submodules=self._sandwich_submodules(), 

116 ), 

117 "ln_final": RMSNormalizationBridge(name="transformer.ln_f", config=self.cfg), 

118 "unembed": UnembeddingBridge(name="lm_head"), 

119 } 

120 

121 def _sandwich_submodules(self) -> dict[str, Any]: 

122 """Build a fresh set of SandwichBlock submodule bridges. 

123 

124 Returns new instances on every call so each of the three block lists 

125 wraps its own live HF modules rather than sharing bridge objects. 

126 

127 Submodule keys mirror the HF attribute names (``norm_1``..``norm_4``, 

128 ``attn``, ``mlp``) so weight-key translation is identity. Attention is 

129 native (combined ``Wqkv`` + ``qk_bias``, RoPE, custom SDPA path), so it 

130 is delegated via ``maintain_native_attention``; only the combined 

131 ``qkv`` and output ``o`` projections are exposed. The gated MLP is 

132 likewise delegated with its combined gate+up ``fc`` and output ``proj``. 

133 """ 

134 attn = AttentionBridge( 

135 name="attn", 

136 config=self.cfg, 

137 submodules={ 

138 "qkv": LinearBridge(name="Wqkv"), 

139 "o": LinearBridge(name="proj"), 

140 }, 

141 maintain_native_attention=True, 

142 requires_attention_mask=True, 

143 # Combined Wqkv projection — no q/k/v submodules to resolve against. 

144 fused_qkv=True, 

145 ) 

146 return { 

147 "norm_1": RMSNormalizationBridge(name="norm_1", config=self.cfg), 

148 "attn": attn, 

149 "norm_2": RMSNormalizationBridge(name="norm_2", config=self.cfg), 

150 "norm_3": RMSNormalizationBridge(name="norm_3", config=self.cfg), 

151 "mlp": self._ungated_mlp(up="fc", down="proj"), 

152 "norm_4": RMSNormalizationBridge(name="norm_4", config=self.cfg), 

153 } 

154 

155 def prepare_loading(self, model_name: str, model_kwargs: dict) -> None: 

156 """Patch Huginn's remote code for transformers v5 compatibility. 

157 

158 Huginn's modeling code targets transformers 4.44; two things break under 

159 v5 (5.8.1), so two patches: 

160 

161 1. Tied-weights format. ``RavenForCausalLM._tied_weights_keys`` is a list 

162 (``["lm_head.weight"]``, the 4.x format), but v5's ``tie_weights`` -> 

163 ``get_expanded_tied_weights_keys`` calls ``.keys()`` on it and raises 

164 ``AttributeError``. The model does not even construct. Rewrite it to 

165 the v5 dict form ``{"lm_head.weight": "transformer.wte.weight"}`` 

166 (Huginn ties ``lm_head`` to ``transformer.wte``). 

167 

168 2. Weight re-init. Under v5's meta-device load-then-materialise flow, 

169 ``PreTrainedModel._init_weights`` is invoked on modules that already 

170 hold checkpoint weights, re-randomising them. Guard it to skip modules 

171 whose parameters are already on a real (non-meta) device — the same 

172 defensive patch openelm.py applies. 

173 

174 Args: 

175 model_name: The HuggingFace model name/path. 

176 model_kwargs: The kwargs dict for from_pretrained(). 

177 """ 

178 # Force-import the modeling module so it appears in sys.modules to patch. 

179 if force_import_remote_class(model_name, "raven_modeling_minimal.RavenForCausalLM") is None: 179 ↛ 180line 179 didn't jump to line 180 because the condition on line 179 was never true

180 return 

181 

182 # Each checkpoint revision gets its own module in sys.modules; patch all. 

183 for module in iter_remote_modeling_modules("raven"): 

184 # Patch 1: tied-weights keys list -> v5 dict form (Huginn ties 

185 # lm_head to transformer.wte). 

186 retie_weights_keys_v5( 

187 getattr(module, "RavenForCausalLM", None), 

188 {"lm_head.weight": "transformer.wte.weight"}, 

189 ) 

190 # Patch 2: don't re-randomise already-loaded weights. 

191 pretrained_class = getattr(module, "RavenPreTrainedModel", None) 

192 if pretrained_class is not None: 192 ↛ 183line 192 didn't jump to line 183 because the condition on line 192 was always true

193 patch_init_weights_skip_loaded(pretrained_class)