Coverage for transformer_lens/model_bridge/supported_architectures/olmo_hybrid.py: 90%
35 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-09-21 19:27 +0000
1"""OLMo Hybrid architecture adapter.
3AllenAI's OLMo Hybrid (``OlmoHybridForCausalLM``, Olmo-Hybrid-7B):
4alternating layer types — OLMo2-style full-attention layers (post-norms
5in the residual, full-width QK-norm, NoPE mode when position embeddings
6are withheld) and GatedDeltaNet linear-attention layers (pre-norm, with
7separate q/k/v short convolutions). Attention stays HF-native; the
8OlmoHybrid GatedDeltaNet variant differs from Qwen3Next's (separate
9q/k/v conv states), so it is delegated opaquely rather than through
10GatedDeltaNetBridge's reimplementation. Generation uses the model's own
11OlmoHybridDynamicCache.
12"""
14from typing import Any
16import torch.nn as nn
18from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter
19from transformer_lens.model_bridge.generalized_components import (
20 AttentionBridge,
21 BlockBridge,
22 EmbeddingBridge,
23 LinearBridge,
24 RMSNormalizationBridge,
25 UnembeddingBridge,
26)
27from transformer_lens.model_bridge.generalized_components.base import (
28 GeneralizedComponent,
29)
32class _OlmoHybridBlockBridge(BlockBridge):
33 """BlockBridge with per-layer-type hook aliases and no hook_resid_mid.
35 hook_resid_mid: no single target fits both layer types (ln2.hook_in is the
36 mid-point on linear-attention layers but the raw attn-branch output on
37 full-attention layers); dropped type-visibly, as on ParallelBlockBridge.
39 hook_attn_out / hook_mlp_out must expose the tensor added to the residual
40 stream, which also differs by layer type: full-attention layers are OLMo2
41 post-norm (contribution = norm output), linear-attention layers are pre-norm
42 (contribution = raw sublayer output). Both candidate targets exist on both
43 layer types, so alias fallback lists cannot discriminate; instead the
44 aliases are selected per layer at bind time in set_original_component.
45 """
47 def __init__(self, *args: Any, **kwargs: Any):
48 super().__init__(*args, **kwargs)
49 if self.hook_aliases is BlockBridge.hook_aliases: 49 ↛ 50line 49 didn't jump to line 50 because the condition on line 49 was never true
50 self.hook_aliases = dict(self.hook_aliases)
51 self.hook_aliases.pop("hook_resid_mid", None)
52 # Full-attention (OLMo2 post-norm) layers route the MLP output through
53 # post_feedforward_layernorm, so the residual-facing MLP output is
54 # ln2_post.hook_out there; linear-attention layers have no ln2_post and
55 # fall back to the raw mlp.hook_out.
56 self.hook_aliases["hook_mlp_out"] = ["ln2_post.hook_out", "mlp.hook_out"]
58 def set_original_component(self, original_component: nn.Module) -> None:
59 super().set_original_component(original_component)
60 if self.hook_aliases is BlockBridge.hook_aliases: 60 ↛ 61line 60 didn't jump to line 61 because the condition on line 60 was never true
61 self.hook_aliases = dict(self.hook_aliases)
62 if getattr(original_component, "post_feedforward_layernorm", None) is not None:
63 # Full-attention (post-norm) layer: ln2 = post_attention_layernorm
64 # applied after attention, ln2_post = post_feedforward_layernorm.
65 # The MLP consumes the mid-residual directly, so the hook_mlp_in
66 # capture must sit on the MLP, not on ln2 (whose input here is the
67 # raw attention output).
68 self.hook_aliases["hook_attn_out"] = "ln2.hook_out"
69 self.hook_aliases["hook_mlp_out"] = "ln2_post.hook_out"
70 self.mlp_reads_resid_directly = True
71 else:
72 # Linear-attention (pre-norm) layer.
73 self.hook_aliases["hook_attn_out"] = "linear_attn.hook_out"
74 self.hook_aliases["hook_mlp_out"] = "mlp.hook_out"
75 self.mlp_reads_resid_directly = False
78class OlmoHybridArchitectureAdapter(ArchitectureAdapter):
79 """Architecture adapter for OlmoHybridForCausalLM models."""
81 # Post-norm attention layers and the linear-attention state are not
82 # fold-safe; compatibility-mode weight processing does not apply.
83 supports_fold_ln = False
85 def __init__(self, cfg: Any) -> None:
86 """Initialize the OLMo Hybrid architecture adapter."""
87 super().__init__(cfg)
89 self._set_rms_rotary_defaults()
90 self.cfg.attn_implementation = "eager"
91 self.cfg.is_stateful = True
93 self.weight_processing_conversions = {
94 **self._qkvo_weight_conversions(),
95 }
97 self.component_mapping = {
98 "embed": EmbeddingBridge(name="model.embed_tokens"),
99 "blocks": _OlmoHybridBlockBridge(
100 name="model.layers",
101 submodules={
102 # Linear-attention layers are pre-norm (input_layernorm);
103 # full-attention layers are OLMo2 post-norm and have
104 # post_feedforward_layernorm instead.
105 "ln1": RMSNormalizationBridge(
106 name="input_layernorm", config=self.cfg, optional=True
107 ),
108 "ln2": RMSNormalizationBridge(name="post_attention_layernorm", config=self.cfg),
109 "ln2_post": RMSNormalizationBridge(
110 name="post_feedforward_layernorm", config=self.cfg, optional=True
111 ),
112 "attn": AttentionBridge(
113 name="self_attn",
114 config=self.cfg,
115 submodules={
116 "q": LinearBridge(name="q_proj"),
117 "k": LinearBridge(name="k_proj"),
118 "v": LinearBridge(name="v_proj"),
119 "o": LinearBridge(name="o_proj"),
120 "q_norm": RMSNormalizationBridge(name="q_norm", config=self.cfg),
121 "k_norm": RMSNormalizationBridge(name="k_norm", config=self.cfg),
122 },
123 maintain_native_attention=True,
124 requires_attention_mask=True,
125 optional=True,
126 ),
127 "linear_attn": GeneralizedComponent(name="linear_attn", optional=True),
128 "mlp": self._gated_mlp(),
129 },
130 ),
131 "ln_final": RMSNormalizationBridge(name="model.norm", config=self.cfg),
132 "unembed": UnembeddingBridge(name="lm_head", config=self.cfg),
133 }
135 def create_stateful_cache(
136 self,
137 hf_model: Any,
138 batch_size: int,
139 device: Any,
140 dtype: Any,
141 ) -> Any:
142 """OLMo Hybrid keeps per-layer q/k/v conv states in its own cache class."""
143 from transformers.models.olmo_hybrid.modeling_olmo_hybrid import (
144 OlmoHybridDynamicCache,
145 )
147 return OlmoHybridDynamicCache(config=hf_model.config)