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

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 

16from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter 

17from transformer_lens.model_bridge.generalized_components import ( 

18 AttentionBridge, 

19 BlockBridge, 

20 EmbeddingBridge, 

21 LinearBridge, 

22 RMSNormalizationBridge, 

23 UnembeddingBridge, 

24) 

25from transformer_lens.model_bridge.generalized_components.base import ( 

26 GeneralizedComponent, 

27) 

28 

29 

30class _OlmoHybridBlockBridge(BlockBridge): 

31 """BlockBridge without the hook_resid_mid alias. 

32 

33 No single target fits both layer types (ln2.hook_in is the mid-point on 

34 linear-attention layers but the raw attn-branch output on full-attention 

35 layers); dropped type-visibly, as on ParallelBlockBridge. 

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 self.hook_aliases.pop("hook_resid_mid", None) 

43 

44 

45class OlmoHybridArchitectureAdapter(ArchitectureAdapter): 

46 """Architecture adapter for OlmoHybridForCausalLM models.""" 

47 

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

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

50 supports_fold_ln = False 

51 

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

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

54 super().__init__(cfg) 

55 

56 self._set_rms_rotary_defaults() 

57 self.cfg.attn_implementation = "eager" 

58 self.cfg.is_stateful = True 

59 

60 self.weight_processing_conversions = { 

61 **self._qkvo_weight_conversions(), 

62 } 

63 

64 self.component_mapping = { 

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

66 "blocks": _OlmoHybridBlockBridge( 

67 name="model.layers", 

68 submodules={ 

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

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

71 # post_feedforward_layernorm instead. 

72 "ln1": RMSNormalizationBridge( 

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

74 ), 

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

76 "ln2_post": RMSNormalizationBridge( 

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

78 ), 

79 "attn": AttentionBridge( 

80 name="self_attn", 

81 config=self.cfg, 

82 submodules={ 

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

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

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

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

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

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

89 }, 

90 maintain_native_attention=True, 

91 requires_attention_mask=True, 

92 optional=True, 

93 ), 

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

95 "mlp": self._gated_mlp(), 

96 }, 

97 ), 

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

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

100 } 

101 

102 def create_stateful_cache( 

103 self, 

104 hf_model: Any, 

105 batch_size: int, 

106 device: Any, 

107 dtype: Any, 

108 ) -> Any: 

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

110 from transformers.models.olmo_hybrid.modeling_olmo_hybrid import ( 

111 OlmoHybridDynamicCache, 

112 ) 

113 

114 return OlmoHybridDynamicCache(config=hf_model.config)