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

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 whose norm keys and hook aliases are chosen per layer type. 

34 

35 ``post_attention_layernorm`` is pre-MLP on linear-attention layers, so only those keep ``ln2``. 

36 """ 

37 

38 def __init__(self, *args: Any, **kwargs: Any): 

39 super().__init__(*args, **kwargs) 

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

41 self.hook_aliases = dict(self.hook_aliases) 

42 # Per-layer targets are set at bind time; this fallback covers the 

43 # template, which has both ln2 and ln2_post declared. 

44 self.hook_aliases["hook_mlp_out"] = ["ln2_post.hook_out", "mlp.hook_out"] 

45 

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

47 super().set_original_component(original_component) 

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

49 self.hook_aliases = dict(self.hook_aliases) 

50 # Runs before setup_submodules binds the norms, so re-keying is safe. 

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

52 if "ln2" in self.submodules: 52 ↛ 54line 52 didn't jump to line 54 because the condition on line 52 was always true

53 self.submodules["ln1_post"] = self.submodules.pop("ln2") 

54 self.hook_aliases["hook_resid_mid"] = "mlp.hook_in" 

55 self.hook_aliases["hook_attn_out"] = "ln1_post.hook_out" 

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

57 self.mlp_reads_resid_directly = True 

58 else: 

59 self.hook_aliases["hook_resid_mid"] = "ln2.hook_in" 

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

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

62 self.mlp_reads_resid_directly = False 

63 

64 

65class OlmoHybridArchitectureAdapter(ArchitectureAdapter): 

66 """Architecture adapter for OlmoHybridForCausalLM models.""" 

67 

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

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

70 supports_fold_ln = False 

71 

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

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

74 super().__init__(cfg) 

75 

76 self._set_rms_rotary_defaults() 

77 self.cfg.attn_implementation = "eager" 

78 self.cfg.is_stateful = True 

79 

80 self.weight_processing_conversions = { 

81 **self._qkvo_weight_conversions(), 

82 } 

83 

84 self.component_mapping = { 

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

86 "blocks": _OlmoHybridBlockBridge( 

87 name="model.layers", 

88 submodules={ 

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

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

91 # post_feedforward_layernorm instead. 

92 "ln1": RMSNormalizationBridge( 

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

94 ), 

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

96 "ln2_post": RMSNormalizationBridge( 

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

98 ), 

99 "attn": AttentionBridge( 

100 name="self_attn", 

101 config=self.cfg, 

102 submodules={ 

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

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

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

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

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

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

109 }, 

110 maintain_native_attention=True, 

111 requires_attention_mask=True, 

112 optional=True, 

113 ), 

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

115 "mlp": self._gated_mlp(), 

116 }, 

117 ), 

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

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

120 } 

121 

122 def create_stateful_cache( 

123 self, 

124 hf_model: Any, 

125 batch_size: int, 

126 device: Any, 

127 dtype: Any, 

128 ) -> Any: 

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

130 from transformers.models.olmo_hybrid.modeling_olmo_hybrid import ( 

131 OlmoHybridDynamicCache, 

132 ) 

133 

134 return OlmoHybridDynamicCache(config=hf_model.config)