Coverage for transformer_lens/model_bridge/supported_architectures/hyenadna.py: 100%

1"""HyenaDNA architecture adapter. 

2 

3HazyResearch/Stanford's HyenaDNA (``HyenaDNAForCausalLM``, remote code): 

4genomic language models built on the Hyena operator — attention-free 

5long-convolution mixing (in_proj -> short conv + implicit modulated 

6filters -> gated multiplicative recombination -> out_proj). Blocks are 

7otherwise transformer-shaped (pre-LN, gelu fc1/fc2 MLP), stacked under 

8``hyena.backbone`` with LayerNorms and single-character DNA vocab. 

9 

10The mixer delegates to HF wholesale (the implicit filter has no 

11attention-shaped reconstruction); in_proj/out_proj are wrapped for 

12hooks. The HF port ships no generate() (not a GenerationMixin), so 

13generation phases are excluded. 

14""" 

15 

16from typing import Any 

17 

18from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter 

19from transformer_lens.model_bridge.generalized_components import ( 

20 BlockBridge, 

21 EmbeddingBridge, 

22 LinearBridge, 

23 NormalizationBridge, 

24 UnembeddingBridge, 

25) 

26from transformer_lens.model_bridge.generalized_components.base import ( 

27 GeneralizedComponent, 

28) 

29 

30 

31class _HyenaBlockBridge(BlockBridge): 

32 """Bare-tensor block whose backbone call mimics a standalone call, so 

33 tuple-normalization is disabled; aliases drop attention (none here) and expose the mixer.""" 

34 

35 hook_aliases = { 

36 "hook_resid_pre": "hook_in", 

37 "hook_resid_mid": "ln2.hook_in", 

38 "hook_resid_post": "hook_out", 

39 "hook_mixer_in": "mixer.hook_in", 

40 "hook_mixer_out": "mixer.hook_out", 

41 "hook_mlp_in": "mlp.in.hook_in", 

42 "hook_mlp_out": "mlp.out.hook_out", 

43 } 

44 

45 @staticmethod 

46 def _is_standalone_hidden_state_call(args: tuple, kwargs: dict) -> bool: 

47 return False 

48 

49 

50class HyenaDNAArchitectureAdapter(ArchitectureAdapter): 

51 """Architecture adapter for HyenaDNAForCausalLM models.""" 

52 

53 applicable_phases: list[int] = [1, 2, 3, 4] 

54 supports_generation: bool = True 

55 # The Hyena FFT convolution has no incremental form (none upstream either); 

56 # generation recomputes the full prefix per step. 

57 supports_kv_cache: bool = False 

58 # The remote forward takes no attention_mask/position_ids kwargs at all. 

59 supports_batched_generation: bool = False 

60 # Attention-free (Hyena convolutions): no q/k/v projections to fold into, 

61 # matching the other delegated-mixer adapters (rwkv, gidd, llada2_moe). 

62 supports_fold_ln: bool = False 

63 

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

65 """Initialize the HyenaDNA architecture adapter.""" 

66 super().__init__(cfg) 

67 

68 self.cfg.normalization_type = "LN" 

69 self.cfg.uses_rms_norm = False 

70 self.cfg.positional_embedding_type = "none" 

71 self.cfg.gated_mlp = False 

72 self.cfg.attn_only = False 

73 self.cfg.final_rms = False 

74 

75 # No attention weights to rearrange. 

76 self.weight_processing_conversions = {} 

77 

78 self.component_mapping = { 

79 "embed": EmbeddingBridge(name="hyena.backbone.embeddings.word_embeddings"), 

80 "blocks": _HyenaBlockBridge( 

81 name="hyena.backbone.layers", 

82 config=self.cfg, 

83 submodules={ 

84 "ln1": NormalizationBridge( 

85 name="norm1", config=self.cfg, use_native_layernorm_autograd=True 

86 ), 

87 # Hyena long-conv mixer: delegated (implicit filters have no 

88 # attention-shaped reconstruction); projections hookable. 

89 "mixer": GeneralizedComponent( 

90 name="mixer", 

91 submodules={ 

92 "in_proj": LinearBridge(name="in_proj"), 

93 "out_proj": LinearBridge(name="out_proj"), 

94 }, 

95 ), 

96 "ln2": NormalizationBridge( 

97 name="norm2", config=self.cfg, use_native_layernorm_autograd=True 

98 ), 

99 "mlp": self._ungated_mlp(up="fc1", down="fc2"), 

100 }, 

101 ), 

102 "ln_final": NormalizationBridge( 

103 name="hyena.backbone.ln_f", config=self.cfg, use_native_layernorm_autograd=True 

104 ), 

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

106 }