--- library_name: transformers tags: - DNA - language-model - StripedHyena - Evo - long-context license: apache-2.0 --- # Evo1-1-7B-131K Minimal HuggingFace port of the **1-7B-131K** variant of [Evo 1](https://huggingface.co/togethercomputer/evo-1-131k-base) -- the ~7B-parameter long-context model trained for sequences up to 131,072 bytes with linearly interpolated RoPE. ## Why this port? [togethercomputer/evo-1-131k-base](https://huggingface.co/togethercomputer/evo-1-131k-base) ships a `trust_remote_code` HF implementation but it has four gaps that force every downstream user to monkey-patch the model: 1. `output_hidden_states=True` is hardcoded to `None` (intermediate embeddings require forward hooks). 2. `output_attentions=True` is unsupported (flash-attn discards the `(B, H, T, T)` matrix; users must patch the attention module). 3. `attn_implementation` cannot be switched at load time - flash_attn is mandatory at every attention layer. 4. The bare backbone is not exposed via `AutoModel.from_pretrained`; only the LM-head wrapper exists. This repo fixes all four. The math is **bit-exact** with the togethercomputer reference (`max_abs_diff = 0.000e+00` at every layer; see [Parity Verification](#parity-verification)). ## Architecture | Parameter | Value | |---|---| | Total parameters | ~7B | | Layers | 32 | | Attention heads | 32 | | Embedding dimension | 4096 | | FFN hidden dimension | 10 928 (gated GELU) | | Vocabulary size | 512 (UTF-8 byte-level) | | Attention layer indices | [8, 16, 24] | | Hyena layer indices | all others | | Hyena state size | 8 | | Positional encoding | RoPE (base = 10000, **linearly scaled by 1/16**) | | Normalization | RMSNorm (eps = 1e-6) | | Architecture | StripedHyena (29 Hyena blocks + 3 causal MHA blocks) | | Max sequence length | 131 072 (training context) | | Training dtype | bfloat16 (Hyena modal-form `poles` / `residues` kept in fp32) | **Vocabulary:** raw UTF-8 byte IDs 0-255; the embedding/output vocabulary is padded to 512 entries. Byte 0 is EOD/EOS and byte 1 is PAD. ## Pretraining - **Objective:** causal byte-level next-token prediction. - **Data:** [OpenGenome](https://huggingface.co/datasets/LongSafari/open-genome), ~300B tokens of prokaryotic whole-genome DNA. The 131k context window was added on top of the 8k base model via continued training with linearly-scaled RoPE. - **Source checkpoint:** `togethercomputer/evo-1-131k-base@1.1_fix`. ## Parity Verification Hidden-state representations verified **bit-exact** (`max_abs_diff = 0.000e+00`) to the togethercomputer reference at all 34 representation levels (token embedding + each of the 32 hybrid blocks + final RMSNorm), using `attn_implementation="flash_attention_2"` in bf16 (matches the reference's backend choice and the trained dtype). Logits from `Evo1ForCausalLM` were also verified bit-exact (top-1 agreement: 128/128 positions). Verified on H100 with PyTorch 2.7.1 / CUDA 12.9. ### Numerical equivalence across attention backends `flash_attention_2` is bit-exact with the original togethercomputer / evo-design implementations (same CUDA kernel). The `sdpa` and `eager` backends use different kernels (PyTorch's bundled flash kernel and pure-PyTorch matmul, respectively); these compute mathematically equivalent attention but accumulate floating-point operations in slightly different orders, producing per-block diffs at the bf16 noise floor (relative error roughly `1e-4` to `1e-2`). Unlike a standard transformer, where attention is softmax-bounded and per-block diffs stay small through the stack, StripedHyena's Hyena layers use an unbounded-gain IIR filter (no softmax) - so any small per-attention-block diff gets amplified by Hyena's filter gain. Across 32 layers this compounds to ~1% relative error in the intermediate residual stream, though the final post-RMSNorm output is bounded. Use `flash_attention_2` if you need to match the reference's activations bit-for-bit. ## Related Models See the full [Evo1 collection](https://huggingface.co/collections/Taykhoom/evo1-6a24ae4a98f04906482db8c2) on the Hub. | Model | Context | Notes | |---|---|---| | [Taykhoom/Evo1-1-7B-8K](https://huggingface.co/Taykhoom/Evo1-1-7B-8K) | 8 192 | Original Evo 1 base model (8k context). | | **[Taykhoom/Evo1-1-7B-131K](https://huggingface.co/Taykhoom/Evo1-1-7B-131K)** | 131 072 | Long-context Evo 1 with linearly-scaled RoPE (131k context). | | [Taykhoom/Evo1-1.5-7B-8K](https://huggingface.co/Taykhoom/Evo1-1.5-7B-8K) | 8 192 | Evo 1.5: Evo 1 (8k) further trained on ~50% more pretraining tokens. | ## Usage > **Note on dtype.** Evo1 was trained in **bfloat16**, with the Hyena `poles` / `residues` (modal-form filter parameters) and rotary `inv_freq` kept in fp32 for numerical stability. `from_pretrained` defaults to bf16 and the bundled remote model code automatically restores these fp32 invariants after checkpoint loading and after later `.to(...)`, `.half()`, or `.bfloat16()` calls; no manual repair call is needed. > **Note on attention backend.** By default, `from_pretrained` selects `attn_implementation="sdpa"` (PyTorch's bundled scaled-dot-product-attention) - this works out of the box without `flash_attn` installed. The original togethercomputer / evo-design implementations use `flash_attn` unconditionally; **for bit-exact reproduction of reference outputs, explicitly pass `attn_implementation="flash_attention_2"`** (and `pip install flash-attn`). See [Numerical equivalence across attention backends](#numerical-equivalence-across-attention-backends) for the magnitude of the difference. ### Embedding generation (no LM head) ```python import torch from transformers import AutoTokenizer, AutoModel tokenizer = AutoTokenizer.from_pretrained("Taykhoom/Evo1-1-7B-131K", trust_remote_code=True) model = AutoModel.from_pretrained( "Taykhoom/Evo1-1-7B-131K", trust_remote_code=True, attn_implementation="flash_attention_2", # strongly recommended for long context ).cuda().eval() seqs = ["ACGTACGTACGT", "GGGTTTAAACCC"] inputs = tokenizer(seqs, return_tensors="pt", padding=True).to(model.device) with torch.no_grad(): out = model(**inputs, output_hidden_states=True) last_hidden = out.last_hidden_state # (B, T, 4096) all_layers = out.hidden_states # tuple of (B, T, 4096), len = 34 (embed + 32 blocks + post-norm) layer_12_emb = all_layers[12] # often used as a "middle" representation ``` ### LM logits ```python import torch from transformers import AutoTokenizer, AutoModelForCausalLM tokenizer = AutoTokenizer.from_pretrained("Taykhoom/Evo1-1-7B-131K", trust_remote_code=True) model = AutoModelForCausalLM.from_pretrained( "Taykhoom/Evo1-1-7B-131K", trust_remote_code=True, attn_implementation="flash_attention_2", ).cuda().eval() inputs = tokenizer(["ACGT"], return_tensors="pt").to(model.device) with torch.no_grad(): logits = model(**inputs).logits # (1, T, 512) ``` ### Faster attention backends Omit `attn_implementation` to use PyTorch SDPA, pass `"eager"` for materialized attention probabilities, or pass `"flash_attention_2"` (with `flash-attn` installed) for bit-exact parity with the upstream implementation. ### Generation ```python import torch from transformers import AutoTokenizer, AutoModelForCausalLM tokenizer = AutoTokenizer.from_pretrained("Taykhoom/Evo1-1-7B-131K", trust_remote_code=True) model = AutoModelForCausalLM.from_pretrained( "Taykhoom/Evo1-1-7B-131K", trust_remote_code=True, attn_implementation="flash_attention_2", ).cuda().eval() inputs = tokenizer(["ACGT"], return_tensors="pt").to(model.device) out = model.generate(**inputs, max_new_tokens=128, do_sample=True, top_k=4, temperature=1.0) print(tokenizer.decode(out[0])) ``` `generation_config.json` ships with `eos_token_id = 0` (the EOD byte) and `pad_token_id = 1` so `model.generate()` stops naturally at the trained end-of-document token without needing extra kwargs. Note that the tokenizer itself does **not** add an EOS at encoding time - this matches the original Evo1 inference pipeline (only generation stops on EOS; embedding/scoring uses raw byte input). > **Cache behavior:** Unpadded single-sequence generation uses the recurrent cache. Padded requests automatically use full recomputation so batched outputs remain correct. Fresh caches consume the complete prompt, and `Evo1Cache.reset()` clears both attention and Hyena state. ### Attention weights ```python import torch from transformers import AutoTokenizer, AutoModel tokenizer = AutoTokenizer.from_pretrained("Taykhoom/Evo1-1-7B-131K", trust_remote_code=True) model = AutoModel.from_pretrained( "Taykhoom/Evo1-1-7B-131K", trust_remote_code=True, attn_implementation="eager", # avoids the automatic eager fallback ).cuda().eval() inputs = tokenizer(["ACGTACGT"], return_tensors="pt").to(model.device) with torch.no_grad(): out = model(**inputs, output_attentions=True) # out.attentions is a tuple of length 32. Entries at indices not in [8, 16, 24] # are None (Hyena blocks have no attention matrix). Entries at [8, 16, 24] are # (B, num_heads, T, T) tensors. attn_block_8 = out.attentions[8] ``` ### Multi-GPU loading (optional) Loading via `accelerate`'s `device_map` is supported (`_no_split_modules` is set so each `AttentionBlock` / `ParallelGatedConvBlock` stays atomic on one device, with hidden state automatically transferred across device boundaries): ```python import torch from transformers import AutoTokenizer, AutoModel tokenizer = AutoTokenizer.from_pretrained("Taykhoom/Evo1-1-7B-131K", trust_remote_code=True) model = AutoModel.from_pretrained( "Taykhoom/Evo1-1-7B-131K", trust_remote_code=True, attn_implementation="flash_attention_2", device_map="auto", # auto-shard across all visible GPUs; falls back to single GPU if only one is present ).eval() ``` Requires `pip install accelerate`. ### Fine-tuning Standard HuggingFace conventions. For sequence-level tasks, take the final `last_hidden_state` (or any intermediate `hidden_states[i]`) and feed it into a downstream head. ## Implementation Notes - **HF base-model contract.** `base_model_prefix = "backbone"` exposes the raw `StripedHyena` module through `.base_model`. Its public `forward` accepts standard HF arguments (`input_ids`, `attention_mask`, `past_key_values`, `use_cache`, `output_hidden_states`, `output_attentions`, `return_dict`) and returns `BaseModelOutputWithPast`; `AutoModel` delegates to this same path. - **Custom attention module (`attention.py`).** Replaces `flash_attn.modules.mha.MHA` with a small in-repo `MHA` class that supports `attn_implementation="eager"` / `"sdpa"` / `"flash_attention_2"`. Learned parameter names (`Wqkv`, `out_proj`) are preserved so existing checkpoints load unchanged. Attention masks are applied to keys in every backend. When `output_attentions=True`, the sdpa and flash paths automatically fall back to eager so the attention matrix is materialized. - **Long-context RoPE (`rotary.py`).** The 131k context window is achieved by linearly scaling RoPE position indices by 1/16 (equivalent to extending the 8k base model's effective context by 16x). Implemented in `LinearlyScaledRotaryEmbedding`, which subclasses either `flash_attn.layers.rotary.RotaryEmbedding` (when available) or our pure-PyTorch fallback. - **Custom rotary embedding (`rotary.py`).** `inv_freq` is treated as non-persistent runtime state and reconstructed from `base` and `dim` in fp32 after checkpoint loading and whenever the RoPE cache is rebuilt. When `flash_attn` is installed we delegate application to its Triton kernel; the pure-PyTorch fallback likewise performs the rotary multiply in fp32 before casting back. This avoids the per-layer error caused by bf16 frequency or multiplication rounding. - **Hyena engine (`engine.py`).** Ported from the togethercomputer reference (FFT-based long convolution, modal-form prefill). - **Cache subclass (`cache.py`).** `Evo1Cache(transformers.cache_utils.Cache)` wraps the two block-type-specific inference param dataclasses (`InferenceParams` for attention KV cache, `RecurrentInferenceParams` for Hyena FIR window + IIR modal state). Exposes `get_seq_length()` / `get_max_cache_shape()` for generation and clears all KV, FIR, and recurrent Hyena state on reset. - **Tokenizer (`tokenization_evo1.py`).** Byte-level UTF-8 with vocab_size = 512. Pad token is byte `\x01`. No CLS, no EOS appended at encoding time (matches original Evo1 inference). Encoding and token-to-ID conversion both preserve byte `\x00` as ID 0; decoding is NumPy-2.x compatible. - **Dependencies.** `torch`, `transformers`, `numpy`, `safetensors`, `huggingface_hub` (only for `from_pretrained` downloads). `flash_attn` is optional: it's used for `attn_implementation="flash_attention_2"` and, when installed, for the RoPE kernel. Without it the model falls back to a pure-PyTorch RoPE with identical results. flash-attn 2.7 and 2.8 both work. ## Citation ```bibtex @article{nguyen2024_evo, title = {Sequence modeling and design from molecular to genome scale with {Evo}}, author = {Nguyen, Eric and Poli, Michael and Durrant, Matthew G. and Kang, Brian and Katrekar, Dhruva and Li, David B. and Bartie, Liam J. and Thomas, Armin W. and King, Samuel H. and Brixi, Garyk and Sullivan, Jeremy and Ng, Madelena Y. and Lewis, Ashley and Lou, Aaron and Ermon, Stefano and Baccus, Stephen A. and Hernandez-Boussard, Tina and {R{\'e}}, Christopher and Hsu, Patrick D. and Hie, Brian L.}, journal = {Science}, volume = {386}, number = {6723}, pages = {eado9336}, year = {2024}, doi = {10.1126/science.ado9336} } ``` ## Credits Original Evo1 model and code by Nguyen et al. Source repo: [evo-design/evo](https://github.com/evo-design/evo). Source checkpoint: [togethercomputer/evo-1-131k-base](https://huggingface.co/togethercomputer/evo-1-131k-base). The HuggingFace conversion code in this repo was authored primarily by [Claude](https://claude.ai/code) and reviewed manually by Taykhoom Dalal. ## License Apache 2.0, following the original Evo1 release.