"""Minimal Evo1 (StripedHyena) HuggingFace port. This module is a refactor of togethercomputer/evo-1-131k-base@1.1_fix's ``modeling_hyena.py`` + ``model.py`` into a single self-contained file with: * ``output_hidden_states`` and ``output_attentions`` plumbed end-to-end, * ``attn_implementation`` switch (``eager`` / ``sdpa`` / ``flash_attention_2``), * ``Evo1Model`` (no LM head, ``BaseModelOutputWithPast``) for ``AutoModel``, * ``Evo1ForCausalLM`` (with logits, ``CausalLMOutputWithPast``) for ``AutoModelForCausalLM``, * minimal external imports (only ``torch`` + ``transformers``; ``flash-attn`` is loaded lazily and only when ``attn_implementation='flash_attention_2'``). Hyena blocks have no attention matrix by construction, so they always emit ``None`` in the per-layer ``attentions`` tuple. Attention blocks (layers 8, 16, 24 for Evo1) emit the (B, H, T, T) softmax matrix when ``output_attentions=True`` (this triggers a one-time fallback from sdpa / flash_attention_2 to the eager backend). """ from __future__ import annotations from typing import Optional, Tuple, Union import torch import torch.nn as nn import torch.nn.functional as F from transformers import PreTrainedModel from transformers.generation import GenerationMixin from transformers.modeling_outputs import ( BaseModelOutputWithPast, CausalLMOutputWithPast, ) from transformers.utils import logging from .attention import MHA from .cache import Evo1Cache, InferenceParams, RecurrentInferenceParams from .configuration_evo1 import Evo1Config from .engine import HyenaInferenceEngine from .layers import ParallelGatedMLP, RMSNorm, VocabParallelEmbedding from .rotary import swap_mha_rope # dummy import so that trust_remote_code bundles the tokenizer file from .tokenization_evo1 import ByteTokenizer # noqa: F401 logger = logging.get_logger(__name__) # ============================================================================= # Block: attention (used at layers config.attn_layer_idxs) # ============================================================================= class AttentionBlock(nn.Module): """Pre-norm Transformer block: norm -> MHA -> residual -> norm -> MLP -> residual.""" def __init__(self, config, layer_idx) -> None: super().__init__() self.config = config self.layer_idx = layer_idx self.pre_norm, self.post_norm = RMSNorm(config), RMSNorm(config) self.proj_groups = config.get("proj_groups", 1) dtype = config.get("attn_block_dtype", torch.bfloat16) mlp_dtype = config.get("mlp_dtype", torch.bfloat16) self.num_attention_heads = config.num_attention_heads self.hidden_size_per_attention_head = ( config.hidden_size // config.num_attention_heads ) attn_impl = getattr(config, "_attn_implementation", "eager") self.inner_mha_cls = MHA( embed_dim=config.hidden_size, num_heads=config.num_attention_heads, num_heads_kv=config.num_attention_heads // self.proj_groups, rotary_emb_dim=config.hidden_size // config.num_attention_heads, qkv_proj_bias=config.get("qkv_proj_bias", True), rotary_emb_base=config.get("rotary_emb_base", 10000), causal=True, layer_idx=layer_idx, out_proj_bias=config.get("mha_out_proj_bias", True), attn_implementation=attn_impl, ).to(dtype=dtype) if config.get("use_interpolated_rotary_pos_emb", False): swap_mha_rope( mha=self.inner_mha_cls, kwargs_new_rope={ "scaling_factor": config.get("rotary_emb_scaling_factor", 1.0) }, ) if self.config.get("smeared_gqa", False): self.inner_mha_cls.num_heads_kv = self.inner_mha_cls.num_heads # Make sure the inv_freq buffer round-trips through to_bfloat16/state_dict. self.inner_mha_cls.rotary_emb.register_buffer( "inv_freq", self.inner_mha_cls.rotary_emb.inv_freq ) self.mlp = ParallelGatedMLP(config).to(dtype=mlp_dtype) def forward( self, u: torch.Tensor, inference_params=None, padding_mask: Optional[torch.Tensor] = None, output_attentions: bool = False, *args, **kwargs, ): if isinstance(padding_mask, torch.Tensor): # Workaround for masking with no qkv bias: this zeros the attended # values at pad positions so they don't leak via attention. u = u * padding_mask[..., None] attn_out, attn_weights = self.inner_mha_cls( self.pre_norm(u), inference_params=inference_params, output_attentions=output_attentions, ) u = attn_out + u if isinstance(padding_mask, torch.Tensor): u = u * padding_mask[..., None] u = self.mlp(self.post_norm(u)) + u return u, attn_weights # ============================================================================= # Block: Hyena (used at all other layers) # ============================================================================= class ParallelHyenaFilter(nn.Module): def __init__(self, config, layer_idx) -> None: super().__init__() self.config = config self.layer_idx = layer_idx self.hyena_filter_groups = config.get( "hyena_filter_groups", self.config.hidden_size ) self.use_flashfft = config.get("use_flashfft", False) self.state_size = config.state_size self.hidden_size = config.hidden_size self.num_filters = config.num_filters self.inference_mode = config.get("inference_mode", True) self.column_split_hyena = config.get("column_split_hyena", True) assert self.hidden_size % self.num_filters == 0 assert self.num_filters <= self.hidden_size self.D = nn.Parameter(torch.zeros(self.hidden_size)) # heads only used to slice post-FIR projections like the checkpoint self.num_attention_heads = config.num_attention_heads self.hidden_size_per_attention_head = ( self.hidden_size // self.num_attention_heads ) self.short_filter_length = config.short_filter_length self.short_filter_weight = nn.Parameter( torch.randn(3 * config.hidden_size, 1, config.short_filter_length) ) self.short_filter_bias = ( nn.Parameter(torch.randn(3 * config.hidden_size)) if config.short_filter_bias else None ) self.engine = HyenaInferenceEngine(layer_idx=layer_idx) self.use_flash_depthwise = config.get("use_flash_depthwise", False) self.data_dtype = None if self.use_flash_depthwise: # importlib avoids the top-level static-import check that HF's # dynamic_module_utils.check_imports performs against the file. import importlib FlashDepthwiseConv1d = importlib.import_module("flashfftconv").FlashDepthwiseConv1d self.fir_fn = FlashDepthwiseConv1d( channels=3 * self.hidden_size, kernel_size=self.short_filter_length, padding=self.short_filter_length - 1, weights=self.short_filter_weight, bias=self.short_filter_bias, device=None, dtype=self.config.get("depthwise_dtype", torch.bfloat16), ) else: self.fir_fn = F.conv1d self.fftconv_fn = None self.long_fir_threshold = config.get("long_fir_threshold", None) if self.long_fir_threshold is not None: assert self.use_flashfft is False, ( "long_fir_threshold not compatible with fused flashfft" ) self.num_systems = self.hidden_size // self.hyena_filter_groups poles = torch.randn(self.num_systems, self.state_size, 1, 2) poles[..., 0] = 1e-2 * torch.randn(self.num_systems, self.state_size, 1) poles[..., 1] = 1e-3 * torch.randn(self.num_systems, self.state_size, 1) self.poles = nn.Parameter(poles) self.residues = nn.Parameter( torch.randn(self.num_systems, self.state_size, 1, 2) ) self.h = None def forward(self, u, inference_params=None, padding_mask=None, *args, **kwargs): if ( inference_params is not None and self.layer_idx in inference_params.fir_state_dict.keys() ): return self.sequential_forward(u, inference_params) return self.parallel_forward(u, inference_params, padding_mask) def parallel_forward(self, u, inference_params=None, padding_mask=None): L = u.shape[1] z_pre, fir_state = self.engine.parallel_fir( self.fir_fn, u, self.short_filter_weight, self.short_filter_bias, L, fir_length=self.short_filter_length, inference_params=inference_params, padding_mask=padding_mask, ) if inference_params: inference_params.fir_state_dict[self.layer_idx] = fir_state if self.h is None: h, _, _, _ = self.compute_filter(L, u.device) else: h = self.h if self.hyena_filter_groups > 1: h = h.repeat_interleave(self.hidden_size // self.hyena_filter_groups, 1) dims = ( self.hidden_size, self.num_attention_heads, self.hidden_size_per_attention_head, self.state_size, self.hyena_filter_groups, ) y = self.engine.parallel_iir( z_pre, h, self.D, L, t=self.t, poles=self.poles, residues=self.residues, dims=dims, inference_params=inference_params, layer_idx=self.layer_idx, prefill_style=self.config.get("prefill_style", "fft"), use_flashfft=self.use_flashfft, fftconv_fn=self.fftconv_fn, column_split_hyena=self.column_split_hyena, long_fir_threshold=self.long_fir_threshold, padding_mask=padding_mask, ) return y, inference_params def sequential_forward(self, u, inference_params): if self.data_dtype is None: self.data_dtype = u.dtype if len(u.shape) > 2: u = u[:, -1] fir_state = inference_params.fir_state_dict[self.layer_idx] iir_state = inference_params.state_dict[self.layer_idx] z_pre, fir_state = self.engine.step_fir( u, fir_state, weight=self.short_filter_weight, bias=self.short_filter_bias, ) if self.column_split_hyena: x_reshaped = z_pre.reshape( z_pre.shape[0], self.num_attention_heads, 3 * self.hidden_size_per_attention_head, ) head = self.hidden_size_per_attention_head x2 = x_reshaped[:, :, :head].reshape(z_pre.shape[0], -1) x1 = x_reshaped[:, :, head : 2 * head].reshape(z_pre.shape[0], -1) v = x_reshaped[:, :, 2 * head:].reshape(z_pre.shape[0], -1) else: x2, x1, v = z_pre.split( [self.hidden_size, self.hidden_size, self.hidden_size], dim=1 ) y, iir_state = self.engine.step_iir( x2, x1, v, self.D, self.residues, self.poles, iir_state, iir_groups=self.hyena_filter_groups, ) inference_params.fir_state_dict[self.layer_idx] = fir_state inference_params.state_dict[self.layer_idx] = iir_state y = y.to(dtype=self.data_dtype) return y[:, None], inference_params def update_time(self, L, device): if not hasattr(self, "t"): self.t = torch.arange(L, device=device)[None, None] elif self.t.shape[-1] < L: self.t = torch.arange(L, device=device)[None, None] else: self.t = self.t[..., :L] def compute_filter(self, L, device): self.update_time(L, device) filter_dtype = torch.float32 residues = torch.view_as_complex(self.residues.to(filter_dtype)) log_poles = torch.view_as_complex(self.poles.to(filter_dtype)).log() h = (residues * (log_poles * self.t).exp()).real.sum(1)[None] return h, filter_dtype, log_poles, residues class ParallelGatedConvBlock(nn.Module): def __init__(self, config, layer_idx) -> None: super().__init__() self.config = config self.layer_idx = layer_idx self.low_mem_mode = config.get("low_mem_mode", False) dtype = config.get("hyena_block_dtype", torch.float32) mlp_dtype = config.get("mlp_dtype", torch.bfloat16) self.pre_norm = RMSNorm(config).to(dtype=dtype) self.post_norm = RMSNorm(config).to(dtype=dtype) self.filter = ParallelHyenaFilter(config, layer_idx).to(dtype=dtype) self.projections = nn.Linear(config.hidden_size, 3 * config.hidden_size) self.out_filter_dense = nn.Linear( config.hidden_size, config.hidden_size ).to(dtype) self.mlp = ParallelGatedMLP(config).to(dtype=mlp_dtype) def forward( self, u, inference_params=None, padding_mask=None, output_attentions: bool = False, *args, **kwargs, ): z = self.projections(self.pre_norm(u)) if isinstance(padding_mask, torch.Tensor): z = z * padding_mask[..., None] z, inference_params = self.filter( z, inference_params=inference_params, padding_mask=padding_mask ) z_in = self.out_filter_dense(z) + u if isinstance(padding_mask, torch.Tensor): z_in = z_in * padding_mask[..., None] y = self.mlp(self.post_norm(z_in)) + z_in # Hyena blocks have no attention matrix. return y, None def get_block(config, layer_idx, flash_fft=None): if layer_idx in config.attn_layer_idxs: return AttentionBlock(config, layer_idx) if layer_idx in config.hyena_layer_idxs: block = ParallelGatedConvBlock(config, layer_idx) if config.get("use_flashfft", False): block.filter.fftconv_fn = flash_fft return block raise NotImplementedError(f"layer_idx {layer_idx} not in attn or hyena indices") # ============================================================================= # Backbone (StripedHyena) # ============================================================================= class StripedHyena(nn.Module): """Pure backbone: token embedding -> N blocks -> RMSNorm. The unembed step is owned by the LM head wrapper, not here, so that ``Evo1Model`` (no LM head) can return the post-norm hidden state as ``last_hidden_state`` cleanly. """ def __init__(self, config): super().__init__() self.config = config self.embedding_layer = VocabParallelEmbedding(config) self.norm = RMSNorm(config) if config.get("final_norm", True) else None if config.get("use_flashfft", False): import importlib FlashFFTConv = importlib.import_module("flashfftconv").FlashFFTConv # NOTE: the original togethercomputer reference had ``config.seqlen`` # here, which is a typo - that attribute doesn't exist on the # config (it's ``max_seqlen``). The bug was unreachable upstream # because ``use_flashfft`` defaults to False; we fix it so the # path is at least loadable for users who do enable it. # FlashFFTConv requires its build-time seqlen to be 2x the # longest input it'll ever see (zero-padding for FFT). self.flash_fft = FlashFFTConv(2 * config.max_seqlen, dtype=torch.bfloat16) else: self.flash_fft = None self.blocks = nn.ModuleList( get_block(config, i, flash_fft=self.flash_fft) for i in range(config.num_layers) ) def forward( self, x: torch.Tensor, inference_params_dict=None, padding_mask: Optional[torch.Tensor] = None, output_hidden_states: bool = False, output_attentions: bool = False, ): x = self.embedding_layer.embed(x) all_hidden_states: list[torch.Tensor] = [] all_attentions: list[Optional[torch.Tensor]] = [] if output_hidden_states: all_hidden_states.append(x) if inference_params_dict is not None: x, inference_params_dict_out = self._stateful_forward( x, inference_params_dict, all_hidden_states=all_hidden_states, all_attentions=all_attentions, output_hidden_states=output_hidden_states, output_attentions=output_attentions, ) else: x, inference_params_dict_out = self._stateless_forward( x, padding_mask=padding_mask, all_hidden_states=all_hidden_states, all_attentions=all_attentions, output_hidden_states=output_hidden_states, output_attentions=output_attentions, ) if self.norm is not None: x = self.norm(x) if output_hidden_states: all_hidden_states.append(x) return x, inference_params_dict_out, all_hidden_states, all_attentions def _stateful_forward( self, x, inference_params_dict, all_hidden_states, all_attentions, output_hidden_states, output_attentions, ): for block_idx, block in enumerate(self.blocks): block_name = ( "mha" if block_idx in self.config.attn_layer_idxs else "hyena" ) inference_params = inference_params_dict[block_name] x, attn = block( x, inference_params=inference_params, output_attentions=output_attentions, ) if output_hidden_states: all_hidden_states.append(x) if output_attentions: all_attentions.append(attn) return x, inference_params_dict def _stateless_forward( self, x, padding_mask, all_hidden_states, all_attentions, output_hidden_states, output_attentions, ): if isinstance(padding_mask, torch.Tensor): x = x * padding_mask[..., None] for block in self.blocks: x, attn = block( x, inference_params=None, padding_mask=padding_mask, output_attentions=output_attentions, ) if output_hidden_states: all_hidden_states.append(x) if output_attentions: all_attentions.append(attn) return x, None def initialize_inference_params(self, max_batch_size: int = 1) -> Evo1Cache: return Evo1Cache( max_seqlen=self.config.get("max_seqlen", 8192), max_batch_size=max_batch_size, short_filter_length=self.config.short_filter_length, state_size=self.config.state_size, ) def to_bfloat16_except_poles_residues(self): """Cast all parameters to bfloat16 except Hyena poles/residues.""" for k, p in self.named_parameters(): if "poles" not in k and "residues" not in k: p.data = p.data.to(torch.bfloat16) # ============================================================================= # HuggingFace wrappers # ============================================================================= class Evo1PreTrainedModel(PreTrainedModel): config_class = Evo1Config base_model_prefix = "backbone" supports_gradient_checkpointing = False _no_split_modules = ["AttentionBlock", "ParallelGatedConvBlock"] _skip_keys_device_placement = "past_key_values" _keys_to_ignore_on_load_missing = [r"freq", r"\.t$"] _keys_to_ignore_on_load_unexpected = [r"fftconv", r"twiddle_factors"] _supports_flash_attn_2 = True _supports_sdpa = True # Hyena filter SSM parameters (poles / residues) MUST stay in fp32: they # parametrize a long-range modal-form filter whose stability collapses # in bf16. HF will keep these in fp32 even when the rest of the model is # loaded in bf16 (or fp16) via the dtype= kwarg of from_pretrained. _keep_in_fp32_modules = ["poles", "residues"] @classmethod def from_pretrained(cls, *args, **kwargs): # Evo1 was trained in bfloat16, with the modal-form filter parameters # (Hyena poles / residues) kept in fp32 via _keep_in_fp32_modules. # bf16 works correctly for all three attention backends (eager, sdpa, # flash_attention_2). Default to bf16 so users don't have to pass it # explicitly; this also silences HF's flash_attention_2 dtype warning # (which inspects the model dtype before force_dtype() runs in __init__). if "dtype" not in kwargs and "torch_dtype" not in kwargs: kwargs["dtype"] = torch.bfloat16 return super().from_pretrained(*args, **kwargs) class Evo1Model(Evo1PreTrainedModel): """Bare backbone: returns ``BaseModelOutputWithPast``. ``last_hidden_state`` is the final (post-RMSNorm) representation, ready to be fed into a downstream head or unembed projection. """ def __init__(self, config: Evo1Config): super().__init__(config) self.backbone = StripedHyena(config) self.config = config self.post_init() self.force_dtype() def force_dtype(self): # Cast everything except poles/residues to bf16 (the trained dtype). # This runs at __init__ time so the model is usable even without an # explicit ``dtype=torch.bfloat16`` kwarg to ``from_pretrained``. self.backbone.to_bfloat16_except_poles_residues() def get_input_embeddings(self): return self.backbone.embedding_layer def set_input_embeddings(self, value): self.backbone.embedding_layer = value def forward( self, input_ids: torch.LongTensor = None, attention_mask: Optional[torch.LongTensor] = None, past_key_values=None, use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, ) -> Union[Tuple, BaseModelOutputWithPast]: return_dict = ( return_dict if return_dict is not None else self.config.use_return_dict ) output_attentions = ( output_attentions if output_attentions is not None else self.config.output_attentions ) output_hidden_states = ( output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states ) # Evo1Model is the bare backbone (no LM head). Default to no caching: # KV caches and Hyena recurrent state are only useful for autoregressive # generation (Evo1ForCausalLM). For embedding extraction the caches # have a large per-layer memory footprint with no benefit. The user # can still opt-in by passing ``use_cache=True`` explicitly. use_cache = use_cache if use_cache is not None else False if use_cache and self.training: use_cache = False inputs = input_ids if use_cache and past_key_values is None: past_key_values = self.backbone.initialize_inference_params( max_batch_size=input_ids.shape[0], ) last_hidden, past_kv, hidden_states, attentions = self.backbone( inputs, padding_mask=attention_mask, inference_params_dict=past_key_values if use_cache else None, output_hidden_states=output_hidden_states, output_attentions=output_attentions, ) if not return_dict: outputs = (last_hidden,) if use_cache: outputs += (past_kv,) if output_hidden_states: outputs += (tuple(hidden_states),) if output_attentions: outputs += (tuple(attentions),) return outputs return BaseModelOutputWithPast( last_hidden_state=last_hidden, past_key_values=past_kv if use_cache else None, hidden_states=tuple(hidden_states) if output_hidden_states else None, attentions=tuple(attentions) if output_attentions else None, ) class Evo1ForCausalLM(Evo1PreTrainedModel, GenerationMixin): """LM head wrapper. Tied to ``backbone.embedding_layer`` (Evo1 ties weights).""" def __init__(self, config: Evo1Config, **kwargs): super().__init__(config, **kwargs) self.backbone = StripedHyena(config) self.config = config # Pad-to-multiple-of for the vocab (matches togethercomputer config). vocab_size = config.vocab_size if vocab_size % config.make_vocab_size_divisible_by != 0: vocab_size += config.make_vocab_size_divisible_by - ( vocab_size % config.make_vocab_size_divisible_by ) self.vocab_size = vocab_size self.post_init() self.force_dtype() def force_dtype(self): self.backbone.to_bfloat16_except_poles_residues() def get_input_embeddings(self): return self.backbone.embedding_layer def set_input_embeddings(self, value): self.backbone.embedding_layer = value def get_output_embeddings(self): return self.backbone.embedding_layer def set_output_embeddings(self, value): self.backbone.embedding_layer = value def forward( self, input_ids: torch.LongTensor = None, attention_mask: Optional[torch.LongTensor] = None, labels: Optional[torch.LongTensor] = None, past_key_values=None, use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, ) -> Union[Tuple, CausalLMOutputWithPast]: return_dict = ( return_dict if return_dict is not None else self.config.use_return_dict ) output_attentions = ( output_attentions if output_attentions is not None else self.config.output_attentions ) output_hidden_states = ( output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states ) use_cache = use_cache if use_cache is not None else self.config.use_cache if use_cache and labels is not None: logger.warning_once( "use_cache=True is incompatible with loss computation; " "disabling cache." ) use_cache = False inputs = input_ids if use_cache: # If the user (or HF generation) didn't pass our Evo1Cache, # initialize a fresh one on the first call. if not isinstance(past_key_values, Evo1Cache): past_key_values = self.backbone.initialize_inference_params( max_batch_size=input_ids.shape[0], ) else: seqlen_offset = past_key_values.seqlen_offset if seqlen_offset == 0: # Prefill done; set offset to prompt length minus the one # token we're about to consume (and that we'll keep # consuming one-at-a-time below). past_key_values.set_offset(input_ids.shape[-1] - 1) else: past_key_values.advance(1) inputs = input_ids[:, -1:] last_hidden, past_kv, hidden_states, attentions = self.backbone( inputs, padding_mask=attention_mask, inference_params_dict=past_key_values if use_cache else None, output_hidden_states=output_hidden_states, output_attentions=output_attentions, ) # Tied unembed: matmul against embedding weights. logits = last_hidden @ self.backbone.embedding_layer.weight.T loss = None if labels is not None: shift_logits = logits[..., :-1, :].contiguous() shift_labels = labels[..., 1:].contiguous() shift_logits = shift_logits.view(-1, self.config.vocab_size) shift_labels = shift_labels.view(-1).to(shift_logits.device) loss = F.cross_entropy(shift_logits, shift_labels) if not return_dict: outputs = (logits,) if use_cache: outputs += (past_kv,) if output_hidden_states: outputs += (tuple(hidden_states),) if output_attentions: outputs += (tuple(attentions),) if loss is not None: outputs = (loss,) + outputs return outputs return CausalLMOutputWithPast( loss=loss, logits=logits, past_key_values=past_kv if use_cache else None, hidden_states=tuple(hidden_states) if output_hidden_states else None, attentions=tuple(attentions) if output_attentions else None, ) @classmethod def can_generate(cls) -> bool: return True def prepare_inputs_for_generation( self, input_ids, attention_mask=None, past_key_values=None, **kwargs ): return { "input_ids": input_ids, "attention_mask": attention_mask, "past_key_values": past_key_values, }