Evo1-1-7B-8K

Evo1-1-7B-8K

Minimal HuggingFace port of the 1-7B-8K variant of Evo 1 -- a ~7B-parameter StripedHyena DNA foundation model trained with an 8,192-token context.

Why this port?

togethercomputer/evo-1-8k-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).

Architecture

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, ~300B tokens of prokaryotic whole-genome DNA.
• Source checkpoint: togethercomputer/[email protected]_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 on the Hub.

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 for the magnitude of the difference.

Memory on long inputs. The depthwise Hyena convolutions and the long-filter construction run 1,024 channels at a time by default (channel_chunk_size). Channels never mix in these operations, so outputs are bit-identical to running all channels at once, but peak memory on long inputs is much lower: on Evo1 7B (131K), a 24 kb input drops from 31 to 23 GiB, and a 205 kb transcript embedded in 131K-token pieces fits in 79 GiB. Pass channel_chunk_size=None to from_pretrained to turn it off.

Embedding generation (no LM head)

import torch
from transformers import AutoTokenizer, AutoModel

tokenizer = AutoTokenizer.from_pretrained("Taykhoom/Evo1-1-7B-8K", trust_remote_code=True)
model = AutoModel.from_pretrained(
    "Taykhoom/Evo1-1-7B-8K",
    trust_remote_code=True,
    attn_implementation="flash_attention_2",  # bit-exact with reference; or omit to default to "sdpa"
).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

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

tokenizer = AutoTokenizer.from_pretrained("Taykhoom/Evo1-1-7B-8K", trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
    "Taykhoom/Evo1-1-7B-8K",
    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

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

tokenizer = AutoTokenizer.from_pretrained("Taykhoom/Evo1-1-7B-8K", trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
    "Taykhoom/Evo1-1-7B-8K",
    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; the prompt is prefilled through the same path as forward_chunked below, bit-identical to a normal forward pass. 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.

Long sequences in chunks

model.backbone.forward_chunked runs a long sequence a chunk at a time, while every position still sees the whole sequence before it. The Hyena blocks carry their short-filter windows and long-filter (IIR) states from chunk to chunk, and the attention blocks a KV cache. Peak activation memory then depends on the chunk size instead of the sequence length; only the KV cache grows with the sequence.

import torch
from transformers import AutoTokenizer, AutoModel

tokenizer = AutoTokenizer.from_pretrained("Taykhoom/Evo1-1-7B-8K", trust_remote_code=True)
model = AutoModel.from_pretrained(
    "Taykhoom/Evo1-1-7B-8K",
    trust_remote_code=True,
    attn_implementation="flash_attention_2",
).cuda().eval()

long_sequence = "ACGT" * 5_000           # 20 kb
ids = tokenizer([long_sequence], return_tensors="pt")["input_ids"].to(model.device)
last_hidden, hidden_states = model.backbone.forward_chunked(
    ids, chunk_size=8192, output_hidden_states=True
)                                      # same outputs as model(ids, output_hidden_states=True)

A single chunk is bit-identical to a normal forward pass. With several chunks, outputs differ from a single pass only by bf16 rounding (in fp32 they agree to about 1e-6), about as much as padding a sequence into a batch changes it. Inputs must be unpadded.

Attention weights

import torch
from transformers import AutoTokenizer, AutoModel

tokenizer = AutoTokenizer.from_pretrained("Taykhoom/Evo1-1-7B-8K", trust_remote_code=True)
model = AutoModel.from_pretrained(
    "Taykhoom/Evo1-1-7B-8K",
    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):

import torch
from transformers import AutoTokenizer, AutoModel

tokenizer = AutoTokenizer.from_pretrained("Taykhoom/Evo1-1-7B-8K", trust_remote_code=True)
model = AutoModel.from_pretrained(
    "Taykhoom/Evo1-1-7B-8K",
    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.
• 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

@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. Source checkpoint: togethercomputer/evo-1-8k-base.

The HuggingFace conversion code in this repo was authored primarily by Claude and reviewed manually by Taykhoom Dalal.

License

Apache 2.0, following the original Evo1 release.