Evo2-1B-8K

A clean, minimal HuggingFace port of Evo 2 1B base, the smallest StripedHyena2 DNA foundation model in the Evo 2 family. Provides native support for layer-by-layer hidden state extraction, attention-weight extraction, and a runtime-switchable attention backend.

NVIDIA Transformer Engine required. This variant uses FP8 input projections (use_fp8_input_projections=True) which require TransformerEngine and a Hopper-class GPU (H100 / H200). Install with:

pip install transformer-engine[pytorch]>=2.3.0

Why this port?

arcinstitute/evo2_1b_base ships a .pt checkpoint that requires the evo2 and vortex Python packages just to instantiate the model. Even with both installed, common pain points remain:

1. Not a HuggingFace model. No from_pretrained, no AutoModel, no AutoModelForCausalLM - the original ships a thin Python wrapper around a custom nn.Module.
2. No public attention-weight API. The default reference path uses Flash Attention and discards the (B, H, T, T) attention matrix.
3. evo2 + vortex packages mandatory even for inference.

This repo fixes all three. Its converted bf16 SDPA path is bit-exact with a matching bf16 SDPA vortex reference configuration (max_abs_diff = 0.000e+00 at every layer; see Parity Verification). Loads with from_pretrained and trust_remote_code=True - no evo2 / vortex install needed.

Architecture

Parameter	Value
Total parameters	~1.1B
Architecture	StripedHyena 2 (interleaved Hyena cascade + MHA blocks)
Layers	25
Attention heads	15
Embedding dimension	1920
Inner MLP size	5 120
Vocabulary size	512 (UTF-8 byte-level)
Attention block indices	3, 10, 17, 24 (4 blocks total)
Hyena block indices	all others (21 blocks: hcs / hcm / hcl pattern)
Positional encoding	RoPE (base = 10 000)
Max sequence length	8 192
Training dtype	bfloat16 (Hyena modal-form log_poles / residues and rotary inv_freq kept in fp32)
FP8 input projections	yes (TransformerEngine required)
Weight format	model.safetensors (1.1B params, single file)

Pretraining

• Objective: causal byte-level next-token prediction.
• Data: OpenGenome2, 8.8 trillion tokens spanning all domains of life.
• Source checkpoint: arcinstitute/evo2_1b_base (evo2_1b_base.pt).

Parity Verification

Hidden-state representations verified bit-exact (max_abs_diff = 0.000e+00) to the vortex reference at every block output after reconstructing the published checkpoint's bf16-rounded inv_freq from base and dim in fp32 on both sides. Parity uses attn_implementation="sdpa" in bf16, with FP8 input projections disabled on both sides where applicable. Logits from Evo2ForCausalLM were also bit-exact (top-1 agreement: 128/128 positions on a 128-byte ACGT input). This verifies conversion fidelity under that controlled configuration; it does not claim bit-exact output against the upstream default Flash Attention / FP8 execution path. Verified on H100 with PyTorch 2.7 / CUDA 12.

Two non-obvious correctness fixes were required versus a naive port (see Implementation Notes for details):

1. inv_freq recomputation. A checkpoint-loaded rotary inv_freq can contain bf16-rounded values even when its destination buffer is fp32. The remote model code treats inv_freq as non-persistent runtime state and reconstructs it in fp32 from base and dim during construction, after checkpoint loading, and whenever the RoPE cache is rebuilt.
2. SDPA backend used for parity. Vortex's reference SelfAttention (use_flash_attn=False) calls F.scaled_dot_product_attention, not a textbook softmax loop. Parity is measured with attn_implementation="sdpa" on our side. Using "eager" (textbook einsum + softmax) is mathematically equivalent but not bit-exact in bf16; using "flash_attention_2" (the recommended runtime backend) is also not bit-exact but agrees within bf16 noise.

Related Models

See the full Evo 2 collection on the Arc Institute HF org for the original weights, or the Taykhoom/Evo2-* collection for our minimal HF ports.

Model	Size	Context	Notes
Taykhoom/Evo2-1B-8K	1B	8 192	<- this model
Taykhoom/Evo2-7B-8K	7B	8 192
Taykhoom/Evo2-7B-262K	7B	262 144
Taykhoom/Evo2-7B-1M	7B	1 048 576
Taykhoom/Evo2-20B-1M	20B	1 048 576
Taykhoom/Evo2-40B-8K	40B	8 192
Taykhoom/Evo2-40B-1M	40B	1 048 576

Usage

Note on dtype. This port supports bfloat16 inference, with the Hyena log_poles / residues (modal-form filter parameters) and rotary inv_freq kept in fp32 for numerical stability. Checkpoint loading normalizes all other parameters to bfloat16 even if another dtype is requested. The bundled remote model code restores the fp32 invariants after loading; no manual repair call is needed.

Note on padding. Use right padding, as shown below. The stateless path zeros padded activations but does not build an additive attention-key mask, so left padding changes logits at otherwise valid positions.

Note on attention backend. By HuggingFace convention this model defaults to attn_implementation="sdpa" (F.scaled_dot_product_attention) since SDPA needs only torch and runs on any GPU. The original Arc Institute Evo 2 inference path uses flash_attention_2, which is faster on long sequences but requires a separate flash-attn install. All usage examples below opt in to flash_attention_2 explicitly because most real users will want it. Drop the kwarg (or pass "sdpa" / "eager") if you don't have flash-attn installed.

Embedding generation (no LM head)

import torch
from transformers import AutoTokenizer, AutoModel

tokenizer = AutoTokenizer.from_pretrained("Taykhoom/Evo2-1B-8K", trust_remote_code=True)
model = AutoModel.from_pretrained(
    "Taykhoom/Evo2-1B-8K",
    trust_remote_code=True,
    attn_implementation="flash_attention_2",  # or "sdpa" (default) or "eager"
).cuda().eval()

seqs = ["ACGTACGTACGT", "GGGTTTAAACCC"]
inputs = tokenizer(seqs, return_tensors="pt", padding=True).to(model.get_input_embeddings().weight.device)

with torch.no_grad():
    out = model(**inputs, output_hidden_states=True)

last_hidden  = out.last_hidden_state   # (B, T, 1920)
all_layers   = out.hidden_states       # tuple of (B, T, 1920), len = 27
middle_layer = all_layers[12]          # input to block 12 (= output of block 11)

Middle-block pre-norm embedding

The Evo 2 paper reports that intermediate representations can outperform final-layer representations, but it does not prescribe this exact layer. This port uses the pre-norm output of the middle block as a simple convention; the upstream README instead demonstrates blocks.28.mlp.l3 for the 7B model. For this variant, the middle-block value is blocks[12].pre_norm(hidden_states[12]):

import torch
from transformers import AutoTokenizer, AutoModel

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

inputs = tokenizer(["ACGTACGTACGT"], return_tensors="pt").to(model.get_input_embeddings().weight.device)
with torch.no_grad():
    out = model(**inputs, output_hidden_states=True)
    pre_norm_middle = model.backbone.blocks[12].pre_norm(
        out.hidden_states[12]
    )                                      # (B, T, 1920)

HF has no built-in API for sub-block intermediates like pre-norm outputs (only block outputs via output_hidden_states). The pattern above applies the block's pre_norm submodule directly to the corresponding hidden_states entry; this gives a bit-identical result to registering a forward hook on backbone.blocks[i].pre_norm and is simpler than using PyTorch hooks. Note that it does require running the full forward pass and then re-applying pre_norm, so a forward hook is more efficient if you only need this single intermediate.

LM logits

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

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

inputs = tokenizer(["ACGT"], return_tensors="pt").to(model.get_input_embeddings().weight.device)
with torch.no_grad():
    logits = model(**inputs).logits   # (1, T, 512)

Generation

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

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


inputs = tokenizer(["ACGT"], return_tensors="pt").to(model.get_input_embeddings().weight.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.

The current cache path ignores attention_mask, so cached batched generation requires equal-length prompts; disable caching for right-padded batches with unequal prompt lengths. The cache also preallocates each attention block to max_seqlen: this variant requires about 0.2 GiB of batch-one attention KV storage at its configured context before model weights. The 1M examples therefore disable caching through model.config.use_cache = False. A per-call generate(..., use_cache=False) is currently dropped by prepare_inputs_for_generation and is not sufficient.

Attention weights

import torch
from transformers import AutoTokenizer, AutoModel

tokenizer = AutoTokenizer.from_pretrained("Taykhoom/Evo2-1B-8K", trust_remote_code=True)
model = AutoModel.from_pretrained(
    "Taykhoom/Evo2-1B-8K",
    trust_remote_code=True,
    attn_implementation="eager",  # required for output_attentions to populate
).cuda().eval()

inputs = tokenizer(["ACGTACGT"], return_tensors="pt").to(model.get_input_embeddings().weight.device)
with torch.no_grad():
    out = model(**inputs, output_attentions=True)

# out.attentions is a tuple of length 25. Entries at indices not in
# [3, 10, 17, 24] are None (Hyena blocks have no attention matrix).
# The 4 attention block(s) at those indices return a (B, num_heads, T, T) tensor.
attn_block_3 = out.attentions[3]

Multi-GPU loading (optional)

For sharding across multiple GPUs (required for 40B, optional for smaller variants), install accelerate and pass device_map="auto":

from transformers import AutoModelForCausalLM
# pip install accelerate
model = AutoModelForCausalLM.from_pretrained(
    "Taykhoom/Evo2-1B-8K", trust_remote_code=True,
    device_map="auto",  # accelerate will shard across all visible GPUs
)

Fine-tuning

This HuggingFace port has not been tested for fine-tuning - it's verified only for inference parity. For fine-tuning, follow the original Arc Institute guidance and use either Savanna (the framework Evo 2 was pretrained in) or NVIDIA BioNeMo, which provides an official Evo 2 fine-tuning recipe.

Implementation Notes

• HF base-model contract. base_model_prefix = "backbone" exposes the raw StripedHyena2 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.
• inv_freq kept exact in fp32 (critical for parity). The geometric series inv_freq[i] = 1 / base^(2i/dim) loses ~7 bits of mantissa when rounded to bf16, which shifts the cos/sin tables and adds Q/K error at every attention layer. inv_freq is non-persistent runtime state: the remote code reconstructs it from base and dim in fp32 after loading and before rebuilding the RoPE cache. Rotary multiplication is also performed in fp32 before casting back to the activation dtype, matching Flash Attention 2.
• log_poles / residues kept in fp32 (critical for stability). The Hyena cascade long (hcl) blocks parameterize an IIR filter via log_poles and residues; bf16 quantisation makes the recurrence numerically unstable. Both are stored as fp32 in the safetensors, covered by _keep_in_fp32_modules, and explicitly restored to fp32 after from_pretrained completes.
• attn_implementation switching (attention.py). Three backends, selected via the standard HF attn_implementation kwarg to from_pretrained (default chosen by HF auto-detection - typically "sdpa"):
  • "sdpa": calls F.scaled_dot_product_attention. Bit-exact with vortex's reference path (when vortex uses use_flash_attn=False).
  • "flash_attention_2": calls flash_attn.flash_attn_qkvpacked_func, the same attention backend used by the default Arc Institute inference path; faster on long sequences; requires flash-attn installed.
  • "eager": textbook einsum + softmax(QK^T) + einsum. Slowest, used internally when output_attentions=True so the attention matrix is materialized.
• Block dispatch (hyena.py). StripedHyena 2 has 4 block types, dispatched by layer_idx membership in four config lists: attn_layer_idxs (MHA + RoPE), hcl_layer_idxs (modal-form IIR via FFT), hcm_layer_idxs (medium FIR cascade, inner length 128), hcs_layer_idxs (short FIR cascade, inner length 7). The disjoint union must equal range(num_layers).
• TELinear with pure-PyTorch fallback (layers.py). Hyena cascade blocks use a TransformerEngine-backed input projection (3x hidden_size output) that supports FP8 quantisation. When TE is not installed, a TELinear fallback class with the same state_dict layout (weight, bias) is used - checkpoints are cross-loadable.
• Custom cache (cache.py). Evo2Cache wraps four block-type-specific dataclasses: InferenceParams for MHA KV cache, HyenaCascadeIIRInferenceParams for hcl, and two HyenaCascadeFIRInferenceParams for hcm / hcs. Passed through model.generate() as past_key_values (we set _supports_cache_class = False so HF treats it as an opaque dict rather than wrapping it in a DynamicCache).
• Tokenizer (tokenization_evo2.py). Byte-level UTF-8, vocab_size = 512. Pad token = byte \x01. EOS = byte \x00 (set as eos_token_id in generation_config.json). Tokenizer does not add EOS at encoding time - matches the original Evo 2 inference pipeline.
• Dependencies. torch, transformers, numpy, safetensors, huggingface_hub. transformer-engine[pytorch] is required for this variant's FP8 input projections. accelerate is optional but recommended if you want to load with device_map="auto" for multi-GPU sharding. flash_attn is optional (only needed if you pass attn_implementation="flash_attention_2").

Citation

@article{brixi2026_evo2,
  title   = {Genome modelling and design across all domains of life with {Evo} 2},
  author  = {Brixi, Garyk and Durrant, Matthew G. and Ku, Jerome and Naghipourfar, Mohsen and Poli, Michael and Sun, Gwanggyu and Brockman, Greg and Chang, Daniel and Fanton, Alison and Gonzalez, Gabriel A. and King, Samuel H. and Li, David B. and Merchant, Aditi T. and Nguyen, Eric and Ricci-Tam, Chiara and Romero, David W. and Schmok, Jonathan C. and Taghibakhshi, Ali and Vorontsov, Anton and Yang, Brandon and Deng, Myra and Gorton, Liv and Nguyen, Nam and Wang, Nicholas K. and Pearce, Michael T. and Simon, Elana and Adams, Etowah and Amador, Zachary J. and Ashley, Euan A. and Baccus, Stephen A. and Dai, Haoyu and Dillmann, Steven and Ermon, Stefano and Guo, Daniel and Herschl, Michael H. and Ilango, Rajesh and Janik, Ken and Lu, Amy X. and Mehta, Reshma and Mofrad, Mohammad R. K. and Ng, Madelena Y. and Pannu, Jaspreet and {R{\'e}}, Christopher and St. John, John and Sullivan, Jeremy and Tey, Joseph and Viggiano, Ben and Zhu, Kevin and Zynda, Greg and Balsam, Daniel and Collison, Patrick and Costa, Anthony B. and Hernandez-Boussard, Tina and Ho, Eric and Liu, Ming-Yu and McGrath, Thomas and Powell, Kimberly and Pinglay, Sudarshan and Burke, Dave P. and Goodarzi, Hani and Hsu, Patrick D. and Hie, Brian L.},
  journal = {Nature},
  volume  = {652},
  number  = {8112},
  pages   = {1349--1361},
  year    = {2026},
  doi     = {10.1038/s41586-026-10176-5}
}

Credits

Original Evo 2 model and code by Brixi et al. (arcinstitute/evo2, Zymrael/vortex). Source checkpoint: arcinstitute/evo2_1b_base.

The HF conversion code was authored primarily by Claude Code and reviewed manually by Taykhoom Dalal.

License

Apache 2.0, following the original Evo 2 release.