Evo2-40B-8K

A clean, minimal HuggingFace port of Evo 2 40B base, the largest 8K-context StripedHyena2 DNA foundation model. 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

Without TE you can load with use_fp8_input_projections=False, which runs the projections in bf16. Outputs then differ noticeably from the original, so use it only when TE isn't an option.

Multi-GPU loading requires accelerate. This variant's bf16 weights (~76 GB) plus activations exceed a single 80 GB H100. Use device_map="auto" to shard across 2 or more H100s; install accelerate first:

pip install accelerate

Why this port?

arcinstitute/evo2_40b_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. With attn_implementation="flash_attention_2" its outputs are bit-exact with the original vortex model (see Parity Verification). Loads with from_pretrained and trust_remote_code=True - no evo2 / vortex install needed.

Architecture

Pretraining

• Objective: causal byte-level next-token prediction.
• Data: OpenGenome2, 8.8 trillion tokens spanning all domains of life.
• Source checkpoint: arcinstitute/evo2_40b_base (evo2_40b_base.pt.part0 and evo2_40b_base.pt.part1; concatenated during conversion).

Parity Verification

With attn_implementation="flash_attention_2", logits and every block output are bit-exact (max_abs_diff = 0.000e+00) against the original vortex model, with the FP8 projections both on and off. We checked this on Evo 2 1B on H100, with PyTorch 2.7 and 2.12 (Transformer Engine 2.5 and 2.17); the other sizes use the same code and checkpoint format. "sdpa" is also bit-exact with vortex's own SDPA path (use_flash_attn=False) when both use the same SDPA backend. PyTorch 2.7 picks its flash backend on H100, while 2.12 picks cuDNN; on newer PyTorch, wrap the forward in torch.nn.attention.sdpa_kernel(SDPBackend.FLASH_ATTENTION) to match. Against the default flash path, "sdpa" and "eager" agree closely but not bit for bit, since they use different attention kernels and the FP8 rounding amplifies the difference.

Two details matter for matching the original exactly (see Implementation Notes):

1. RoPE frequencies. The checkpoint stores inv_freq in bf16, and vortex runs RoPE on those rounded values. The port rebuilds the same rounded values from base and dim. Using the exact values instead changes the frequencies by up to ~0.4%, which is small but enough to break parity.
2. Trained FP8 scales. The original checkpoint carries per-layer FP8 scaling factors. The port loads them and keeps them fixed at inference, as vortex does. Without them Transformer Engine starts every scale at 1.0 and the outputs drift.

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.

Usage

Note on dtype. Bfloat16 is recommended. Float32 is also supported, while float16 is rejected because the modal filters are numerically unstable in FP16. Hyena log_poles / residues and rotary inv_freq remain fp32 for numerical stability.

Note on padding. Attention masks accept boolean, integer, or floating-point 1/0 values. Padded keys are excluded from attention.

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.

Optional Triton kernels. For faster long-sequence inference, pass use_hcs_kernel=True, use_hcm_kernel=True, use_hcl_kernel=True to from_pretrained (requires triton). These are the Hyena kernels from vortex PR #77 by @AlphaKhaw, vendored unchanged in kernels.py. On Evo 2 7B (1M) they cut peak memory by ~40% at 64K tokens, fit 131K tokens on a single 80 GB H100, and run ~1.2x faster. They are inference-only and not bit-exact with the default path (top-token agreement >= 99.98% in our tests); upstream notes they regress generation quality on the 1B.

Embedding generation (no LM head)

import torch
from transformers import AutoTokenizer, AutoModel

tokenizer = AutoTokenizer.from_pretrained("Taykhoom/Evo2-40B-8K", trust_remote_code=True)
model = AutoModel.from_pretrained(
    "Taykhoom/Evo2-40B-8K",
    trust_remote_code=True,
    attn_implementation="flash_attention_2",  # or "sdpa" (default) or "eager"
    device_map="auto",
).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, 8192)
all_layers   = out.hidden_states       # tuple of (B, T, 8192), len = 52
middle_layer = all_layers[25]          # input to block 25 (= output of block 24)

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[25].pre_norm(hidden_states[25]):

import torch
from transformers import AutoTokenizer, AutoModel

tokenizer = AutoTokenizer.from_pretrained("Taykhoom/Evo2-40B-8K", trust_remote_code=True)
model = AutoModel.from_pretrained(
    "Taykhoom/Evo2-40B-8K",
    trust_remote_code=True,
    attn_implementation="flash_attention_2",
    device_map="auto",
).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[25].pre_norm(
        out.hidden_states[25]
    )                                      # (B, T, 8192)

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-40B-8K", trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
    "Taykhoom/Evo2-40B-8K", trust_remote_code=True,
    attn_implementation="flash_attention_2",
    device_map="auto",
).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-40B-8K", trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
    "Taykhoom/Evo2-40B-8K", trust_remote_code=True,
    attn_implementation="flash_attention_2",
    device_map="auto",
).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.

Generation uses recurrent MHA/Hyena state for unpadded prompts. Padded batches automatically use the exact full-recomputation path. Attention KV storage grows with the decoded sequence instead of preallocating the model's full configured context.

Attention weights

import torch
from transformers import AutoTokenizer, AutoModel

tokenizer = AutoTokenizer.from_pretrained("Taykhoom/Evo2-40B-8K", trust_remote_code=True)
model = AutoModel.from_pretrained(
    "Taykhoom/Evo2-40B-8K",
    trust_remote_code=True,
    attn_implementation="eager",  # required for output_attentions to populate
    device_map="auto",
).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 50. Entries at indices not in
# [3, 10, 17, 24, 31, 35, 42, 49] are None (Hyena blocks have no attention matrix).
# The 8 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-40B-8K", trust_remote_code=True,
    device_map="auto",  # accelerate will shard across all visible GPUs
)

Note on training across GPUs. The FP8 projections use the same Transformer Engine recipe as the original Evo 2, which can't be sharded across GPUs when gradients are on. Inference is unaffected. To fine-tune on multiple GPUs, load with fp8_recipe="current"; its outputs differ slightly from the original's, so keep the default for inference.

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.
• RoPE inv_freq (critical for parity). inv_freq is non-persistent runtime state: the remote code rebuilds it from base and dim after loading and whenever the RoPE cache is rebuilt. The values are rounded through bf16 to match the checkpoint and the original vortex runtime, then held in fp32 so HF dtype casts leave them alone. 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 use_flash_attn=False path when both use the same SDPA backend (see Parity Verification).
  • "flash_attention_2": calls flash_attn.flash_attn_qkvpacked_func, the same attention backend used by the default Arc Institute inference path, and bit-exact with it; 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. The trained FP8 scaling factors are stored as fp8_* buffers and handed to TE before the first FP8 call.
• Generation cache. Unpadded prompts use recurrent MHA KV plus HCL/HCM/HCS FIR/IIR state matching the upstream Vortex cache path. Padded batches use full recomputation. KV tensors grow geometrically with the decoded length instead of allocating max_seqlen up front.
• 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 required if you load with device_map="auto" (the model is too large to fit on a single 80 GB H100 with activations). 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{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_40b_base.

Hugging Face port maintained by Taykhoom Dalal.

License

Apache 2.0, following the original Evo 2 release.