Coverage for transformer_lens/model_bridge/supported_architectures/olmo_hybrid.py: 90%

1"""OLMo Hybrid architecture adapter. 

2 

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""" 

13 

14from typing import Any 

15 

16import torch.nn as nn 

17 

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) 

30 

31 

32class _OlmoHybridBlockBridge(BlockBridge): 

33 """BlockBridge with per-layer-type hook aliases and no hook_resid_mid. 

34 

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. 

38 

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 """ 

46 

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 

53 def set_original_component(self, original_component: nn.Module) -> None: 

54 super().set_original_component(original_component) 

55 if self.hook_aliases is BlockBridge.hook_aliases: 55 ↛ 56line 55 didn't jump to line 56 because the condition on line 55 was never true

56 self.hook_aliases = dict(self.hook_aliases) 

57 if getattr(original_component, "post_feedforward_layernorm", None) is not None: 

58 # Full-attention (post-norm) layer: ln2 = post_attention_layernorm 

59 # applied after attention, ln2_post = post_feedforward_layernorm. 

60 # The MLP consumes the mid-residual directly, so the hook_mlp_in 

61 # capture must sit on the MLP, not on ln2 (whose input here is the 

62 # raw attention output). 

63 self.hook_aliases["hook_attn_out"] = "ln2.hook_out" 

64 self.hook_aliases["hook_mlp_out"] = "ln2_post.hook_out" 

65 self.mlp_reads_resid_directly = True 

66 else: 

67 # Linear-attention (pre-norm) layer. 

68 self.hook_aliases["hook_attn_out"] = "linear_attn.hook_out" 

69 self.hook_aliases["hook_mlp_out"] = "mlp.hook_out" 

70 self.mlp_reads_resid_directly = False 

71 

72 

73class OlmoHybridArchitectureAdapter(ArchitectureAdapter): 

74 """Architecture adapter for OlmoHybridForCausalLM models.""" 

75 

76 # Post-norm attention layers and the linear-attention state are not 

77 # fold-safe; compatibility-mode weight processing does not apply. 

78 supports_fold_ln = False 

79 

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

81 """Initialize the OLMo Hybrid architecture adapter.""" 

82 super().__init__(cfg) 

83 

84 self._set_rms_rotary_defaults() 

85 self.cfg.attn_implementation = "eager" 

86 self.cfg.is_stateful = True 

87 

88 self.weight_processing_conversions = { 

89 **self._qkvo_weight_conversions(), 

90 } 

91 

92 self.component_mapping = { 

93 "embed": EmbeddingBridge(name="model.embed_tokens"), 

94 "blocks": _OlmoHybridBlockBridge( 

95 name="model.layers", 

96 submodules={ 

97 # Linear-attention layers are pre-norm (input_layernorm); 

98 # full-attention layers are OLMo2 post-norm and have 

99 # post_feedforward_layernorm instead. 

100 "ln1": RMSNormalizationBridge( 

101 name="input_layernorm", config=self.cfg, optional=True 

102 ), 

103 "ln2": RMSNormalizationBridge(name="post_attention_layernorm", config=self.cfg), 

104 "ln2_post": RMSNormalizationBridge( 

105 name="post_feedforward_layernorm", config=self.cfg, optional=True 

106 ), 

107 "attn": AttentionBridge( 

108 name="self_attn", 

109 config=self.cfg, 

110 submodules={ 

111 "q": LinearBridge(name="q_proj"), 

112 "k": LinearBridge(name="k_proj"), 

113 "v": LinearBridge(name="v_proj"), 

114 "o": LinearBridge(name="o_proj"), 

115 "q_norm": RMSNormalizationBridge(name="q_norm", config=self.cfg), 

116 "k_norm": RMSNormalizationBridge(name="k_norm", config=self.cfg), 

117 }, 

118 maintain_native_attention=True, 

119 requires_attention_mask=True, 

120 optional=True, 

121 ), 

122 "linear_attn": GeneralizedComponent(name="linear_attn", optional=True), 

123 "mlp": self._gated_mlp(), 

124 }, 

125 ), 

126 "ln_final": RMSNormalizationBridge(name="model.norm", config=self.cfg), 

127 "unembed": UnembeddingBridge(name="lm_head", config=self.cfg), 

128 } 

129 

130 def create_stateful_cache( 

131 self, 

132 hf_model: Any, 

133 batch_size: int, 

134 device: Any, 

135 dtype: Any, 

136 ) -> Any: 

137 """OLMo Hybrid keeps per-layer q/k/v conv states in its own cache class.""" 

138 from transformers.models.olmo_hybrid.modeling_olmo_hybrid import ( 

139 OlmoHybridDynamicCache, 

140 ) 

141 

142 return OlmoHybridDynamicCache(config=hf_model.config)